"""
Model architectures for Malicious URL Detection — GPU-accelerated.

Provides:
  • build_xgb(n)       – XGBoost classifier (GPU via device="cuda")
  • build_lgbm(n)      – LightGBM classifier (CPU — GPU build unreliable)
  • build_tfidf_lr()   – TF-IDF char-ngram + Logistic Regression pipeline
  • DeepURLClassifier  – sklearn-compatible wrapper for CharCNN / BiLSTM / CNN-BiLSTM
  • ManualStackingEnsemble – OOF stacking with XGBoost GPU meta-learner
"""

from __future__ import annotations

import warnings
import numpy as np
import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler
try:
    from torch.amp import autocast          # PyTorch >= 2.4
except ImportError:
    from torch.cuda.amp import autocast      # PyTorch < 2.4
from sklearn.base import BaseEstimator, ClassifierMixin, clone
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.model_selection import StratifiedKFold

# ═════════════════════════════════════════════════════════════════════════════
# 1. Character encoding (vectorised, numpy)
# ═════════════════════════════════════════════════════════════════════════════

MAX_URL_LEN = 200

# Build lookup table:  ord(char) → token_id
_CHAR_LUT = np.zeros(256, dtype=np.int64)   # covers full byte range
_PRINTABLE = (
    "abcdefghijklmnopqrstuvwxyz"
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "0123456789"
    "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~ "
)
for _i, _ch in enumerate(_PRINTABLE, start=1):
    _CHAR_LUT[ord(_ch)] = _i
_UNK_ID = len(_PRINTABLE) + 1   # 96
VOCAB_SIZE = _UNK_ID + 1        # 97  (indices 0..96)


def url_to_tensor(urls: list[str], max_len: int = MAX_URL_LEN) -> torch.Tensor:
    """Vectorised char→int encoding for a batch of URL strings.

    Returns a LongTensor of shape ``(len(urls), max_len)``.
    """
    arr = np.zeros((len(urls), max_len), dtype=np.int64)
    for i, url in enumerate(urls):
        raw = str(url)[:max_len]
        chars = np.frombuffer(raw.encode("ascii", "ignore"), dtype=np.uint8)
        if len(chars) == 0:
            continue
        ids = _CHAR_LUT[chars]                       # mapped chars → id, unmapped → 0
        ids = np.where(ids > 0, ids, _UNK_ID)        # replace 0 (unmapped) with UNK
        ids[0 : min(len(ids), max_len)]               # safety slice
        arr[i, : len(ids)] = ids[: max_len]
    return torch.from_numpy(arr)


# ═════════════════════════════════════════════════════════════════════════════
# 2. Deep Learning architectures
# ═════════════════════════════════════════════════════════════════════════════

class CharCNNNet(nn.Module):
    """1-D CNN on character-level URL embeddings."""

    def __init__(self, num_classes: int = 2, embed_dim: int = 64):
        super().__init__()
        self.embed = nn.Embedding(VOCAB_SIZE, embed_dim, padding_idx=0)
        self.convs = nn.Sequential(
            nn.Conv1d(embed_dim, 128, kernel_size=3, padding=1),
            nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2),
            nn.Conv1d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2),
            nn.Conv1d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm1d(256), nn.ReLU(), nn.AdaptiveMaxPool1d(1),
        )
        self.head = nn.Sequential(
            nn.Dropout(0.3), nn.Linear(256, 128), nn.ReLU(),
            nn.Dropout(0.2), nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.embed(x).permute(0, 2, 1)  # (B, C, L)
        x = self.convs(x).squeeze(-1)       # (B, 256)
        return self.head(x)


class CharBiLSTMNet(nn.Module):
    """Bidirectional LSTM on character embeddings."""

    def __init__(self, num_classes: int = 2, embed_dim: int = 64,
                 hidden: int = 128, n_layers: int = 2):
        super().__init__()
        self.embed = nn.Embedding(VOCAB_SIZE, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            embed_dim, hidden, num_layers=n_layers,
            batch_first=True, bidirectional=True, dropout=0.3,
        )
        self.head = nn.Sequential(
            nn.Dropout(0.3), nn.Linear(hidden * 2, 128), nn.ReLU(),
            nn.Dropout(0.2), nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.embed(x)                       # (B, L, E)
        _, (h, _) = self.lstm(x)                 # h: (layers*2, B, H)
        h = torch.cat([h[-2], h[-1]], dim=-1)    # (B, H*2)
        return self.head(h)


class CNNBiLSTMNet(nn.Module):
    """Hybrid: CNN feature extractor → BiLSTM sequence modelling."""

    def __init__(self, num_classes: int = 2, embed_dim: int = 64, hidden: int = 128):
        super().__init__()
        self.embed = nn.Embedding(VOCAB_SIZE, embed_dim, padding_idx=0)
        self.cnn = nn.Sequential(
            nn.Conv1d(embed_dim, 128, 3, padding=1),
            nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2),
        )
        self.lstm = nn.LSTM(
            128, hidden, num_layers=1,
            batch_first=True, bidirectional=True, dropout=0.0,
        )
        self.head = nn.Sequential(
            nn.Dropout(0.3), nn.Linear(hidden * 2, 128), nn.ReLU(),
            nn.Dropout(0.2), nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.embed(x).permute(0, 2, 1)       # (B, C, L)
        x = self.cnn(x).permute(0, 2, 1)          # (B, L', 128)
        _, (h, _) = self.lstm(x)
        h = torch.cat([h[-2], h[-1]], dim=-1)
        return self.head(h)


_ARCH_MAP = {
    "cnn": CharCNNNet,
    "bilstm": CharBiLSTMNet,
    "cnn_bilstm": CNNBiLSTMNet,
}


# ═════════════════════════════════════════════════════════════════════════════
# 3. sklearn-compatible deep learning wrapper (GPU + AMP)
# ═════════════════════════════════════════════════════════════════════════════

class DeepURLClassifier(BaseEstimator, ClassifierMixin):
    """scikit-learn interface for character-level DL URL models.

    Features:
      • Automatic mixed precision (AMP) for speed
      • Multi-GPU via ``nn.DataParallel``
      • Batch prediction to avoid OOM
      • torch.compile on single-GPU setups
    """

    def __init__(
        self,
        arch: str = "cnn",
        num_classes: int = 2,
        epochs: int = 20,
        batch_size: int = 2048,
        lr: float = 1e-3,
        device: str = "auto",
        max_len: int = MAX_URL_LEN,
    ):
        self.arch = arch
        self.num_classes = num_classes
        self.epochs = epochs
        self.batch_size = batch_size
        self.lr = lr
        self.device = device
        self.max_len = max_len

    def _resolve_device(self):
        if self.device == "auto":
            return "cuda" if torch.cuda.is_available() else "cpu"
        return self.device

    def fit(self, X, y, class_weight=None):
        dev = self._resolve_device()
        net = _ARCH_MAP[self.arch](num_classes=self.num_classes).to(dev)

        # Multi-GPU
        n_gpus = torch.cuda.device_count()
        if dev == "cuda" and n_gpus > 1:
            net = nn.DataParallel(net)
            print(f"    DataParallel across {n_gpus} GPUs")

        self.net_ = net

        # Loss with optional class weight
        weight = None
        if class_weight == "balanced":
            counts = np.bincount(np.asarray(y).astype(int), minlength=self.num_classes)
            w = len(y) / (self.num_classes * np.maximum(counts, 1))
            weight = torch.tensor(w, dtype=torch.float32, device=dev)

        criterion = nn.CrossEntropyLoss(weight=weight)
        optimiser = torch.optim.AdamW(net.parameters(), lr=self.lr, weight_decay=1e-4)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=self.epochs)
        _use_amp = (dev == "cuda")
        scaler = GradScaler(enabled=_use_amp)

        urls = list(X)
        labels = torch.tensor(np.asarray(y).astype(np.int64), device=dev)

        net.train()
        for epoch in range(self.epochs):
            perm = torch.randperm(len(urls))
            total_loss = 0.0
            n_batches = 0
            for start in range(0, len(urls), self.batch_size):
                idx = perm[start: start + self.batch_size]
                batch_urls = [urls[i] for i in idx.tolist()]
                xb = url_to_tensor(batch_urls, self.max_len).to(dev)
                yb = labels[idx]

                optimiser.zero_grad(set_to_none=True)
                with autocast("cuda", enabled=_use_amp):
                    logits = net(xb)
                    loss = criterion(logits, yb)
                scaler.scale(loss).backward()
                scaler.step(optimiser)
                scaler.update()
                total_loss += loss.item()
                n_batches += 1
            scheduler.step()
            if (epoch + 1) % 5 == 0 or epoch == 0:
                print(f"    Epoch {epoch+1:>3d}/{self.epochs}  loss={total_loss/n_batches:.4f}")
        return self

    @torch.inference_mode()
    def predict_proba(self, X):
        dev = self._resolve_device()
        net = self.net_
        net.eval()
        _use_amp = (dev == "cuda")
        urls = list(X)
        all_probs = []
        for start in range(0, len(urls), self.batch_size):
            batch = urls[start: start + self.batch_size]
            xb = url_to_tensor(batch, self.max_len).to(dev)
            with autocast("cuda", enabled=_use_amp):
                logits = net(xb)
            probs = torch.softmax(logits.float(), dim=-1).cpu().numpy()
            all_probs.append(probs)
        return np.concatenate(all_probs, axis=0)

    def predict(self, X):
        return np.argmax(self.predict_proba(X), axis=1)


# ═════════════════════════════════════════════════════════════════════════════
# 4. Classical ML builders
# ═════════════════════════════════════════════════════════════════════════════

def build_xgb(num_classes: int = 2, **extra):
    """XGBoost classifier with GPU acceleration."""
    import xgboost as xgb

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    objective = "binary:logistic" if num_classes == 2 else "multi:softprob"
    params = dict(
        n_estimators=500,
        max_depth=8,
        learning_rate=0.05,
        tree_method="hist",
        device=dev,
        objective=objective,
        eval_metric="logloss" if num_classes == 2 else "mlogloss",
        n_jobs=-1,
        random_state=42,
    )
    if num_classes > 2:
        params["num_class"] = num_classes
    params.update(extra)
    return xgb.XGBClassifier(**params)


def build_lgbm(num_classes: int = 2, **extra):
    """LightGBM classifier (CPU only — GPU build unreliable on cloud)."""
    import lightgbm as lgb

    objective = "binary" if num_classes == 2 else "multiclass"
    params = dict(
        n_estimators=500,
        max_depth=8,
        learning_rate=0.05,
        num_leaves=63,
        objective=objective,
        class_weight="balanced",
        device_type="cpu",
        n_jobs=-1,
        random_state=42,
        verbose=-1,
    )
    if num_classes > 2:
        params["num_class"] = num_classes
    params.update(extra)
    return lgb.LGBMClassifier(**params)


def build_tfidf_lr_pipeline(num_classes: int = 2):
    """TF-IDF character n-gram + Logistic Regression pipeline."""
    return Pipeline([
        ("tfidf", TfidfVectorizer(
            analyzer="char_wb",
            ngram_range=(2, 5),
            max_features=100_000,
            sublinear_tf=True,
            dtype=np.float32,
        )),
        ("lr", LogisticRegression(
            C=1.0,
            max_iter=300,
            solver="saga",
            n_jobs=-1,
            random_state=42,
        )),
    ])


# ═════════════════════════════════════════════════════════════════════════════
# 5. Stacking HyperModel
# ═════════════════════════════════════════════════════════════════════════════

def build_stacking_hypermodel(num_classes: int = 2, use_deep: bool = True, cv: int = 5):
    """Return a dict with ``level0`` models and ``meta`` learner config.

    Level-0 models:
      • ``xgb``      – XGBoost (GPU, features)
      • ``lgbm``     – LightGBM (CPU, features)
      • ``char_cnn`` – CharCNN (GPU, URL strings)     [if use_deep]
      • ``char_bilstm`` – CharBiLSTM (GPU, URL strings) [if use_deep]

    Meta-learner: XGBoost (GPU).
    """
    level0 = {
        "xgb": build_xgb(num_classes),
        "lgbm": build_lgbm(num_classes),
    }
    if use_deep:
        level0["char_cnn"] = DeepURLClassifier(
            arch="cnn", num_classes=num_classes, epochs=15,
            batch_size=2048, lr=1e-3, device="auto",
        )
        level0["char_bilstm"] = DeepURLClassifier(
            arch="bilstm", num_classes=num_classes, epochs=12,
            batch_size=2048, lr=5e-4, device="auto",
        )

    meta = build_xgb(num_classes, n_estimators=300, max_depth=4, learning_rate=0.1)
    return {"level0": level0, "meta": meta, "cv": cv}


class ManualStackingEnsemble(BaseEstimator, ClassifierMixin):
    """Out-of-fold stacking ensemble.

    Generates OOF probability predictions from each L0 model across
    ``cv`` folds, stacks them as features, and trains an XGBoost
    meta-learner on those stacked predictions.
    """

    def __init__(self, level0: dict, meta, cv: int = 5, num_classes: int = 2):
        self.level0 = level0
        self.meta = meta
        self.cv = cv
        self.num_classes = num_classes

    def _uses_urls(self, name: str) -> bool:
        """Whether a model operates on raw URL strings."""
        return name in ("tfidf_lr", "char_cnn", "char_bilstm")

    def fit(self, X_feat, X_url, y):
        """Train all L0 models via OOF, then train the meta-learner.

        Parameters
        ----------
        X_feat : np.ndarray  (n, d)   – scaled feature matrix.
        X_url  : np.ndarray  (n,)     – raw URL strings.
        y      : np.ndarray  (n,)     – labels.
        """
        y = np.asarray(y)
        n = len(y)
        oof = np.zeros((n, self.num_classes * len(self.level0)), dtype=np.float32)
        kf = StratifiedKFold(n_splits=self.cv, shuffle=True, random_state=42)

        self.fitted_l0_ = {}

        for model_idx, (name, base) in enumerate(self.level0.items()):
            print(f"\n  [INFO] OOF training: {name}")
            col_start = model_idx * self.num_classes
            col_end = col_start + self.num_classes
            fold_models = []

            for fold, (tr_idx, val_idx) in enumerate(kf.split(X_feat, y)):
                mdl = clone(base)

                if self._uses_urls(name):
                    X_tr = X_url[tr_idx]
                    X_val = X_url[val_idx]
                else:
                    X_tr = X_feat[tr_idx]
                    X_val = X_feat[val_idx]

                y_tr = y[tr_idx]

                # Only DeepURLClassifier gets class_weight
                if isinstance(mdl, DeepURLClassifier):
                    mdl.fit(X_tr, y_tr, class_weight="balanced")
                else:
                    mdl.fit(X_tr, y_tr)

                proba = np.asarray(mdl.predict_proba(X_val))
                oof[val_idx, col_start:col_end] = proba
                fold_models.append(mdl)

            # Retrain on full data for inference
            final = clone(base)
            if self._uses_urls(name):
                X_full = X_url
            else:
                X_full = X_feat
            if isinstance(final, DeepURLClassifier):
                final.fit(X_full, y, class_weight="balanced")
            else:
                final.fit(X_full, y)
            self.fitted_l0_[name] = final

        # Train meta-learner on OOF predictions
        print("\n  [INFO] Training meta-learner ...")
        self.meta.fit(oof, y)
        self.oof_ = oof
        return self

    def _build_meta_features(self, X_feat, X_url):
        """Stack L0 predictions into a meta-feature matrix."""
        parts = []
        for name, mdl in self.fitted_l0_.items():
            if self._uses_urls(name):
                proba = np.asarray(mdl.predict_proba(X_url))
            else:
                proba = np.asarray(mdl.predict_proba(X_feat))
            parts.append(proba)
        return np.hstack(parts).astype(np.float32)

    def predict_proba(self, X_feat, X_url):
        meta_X = self._build_meta_features(X_feat, X_url)
        return np.asarray(self.meta.predict_proba(meta_X))

    def predict(self, X_feat, X_url):
        return np.argmax(self.predict_proba(X_feat, X_url), axis=1)