src/models/deep_learning.py

"""
Deep learning models for user behavior profiling.

Implements LSTM and Transformer architectures that model temporal sequences
of cardholder transactions to detect anomalous behavior patterns indicative
of fraud (account takeover, gradual compromise, etc.).
"""

import logging
import sys
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
from sklearn.metrics import average_precision_score, roc_auc_score
from torch.utils.data import DataLoader, Dataset

logger = logging.getLogger(__name__)

if torch.cuda.is_available():
    DEVICE = torch.device("cuda")
else:
    # MPS (Apple Silicon) has known issues with bidirectional LSTM and some
    # Transformer ops — use CPU for reliability on M1/M2/M3 Macs.
    DEVICE = torch.device("cpu")


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------

class TransactionSequenceDataset(Dataset):
    """Dataset that creates fixed-length transaction sequences per cardholder.

    Each sample is a sequence of the most recent N transactions for a cardholder,
    with the label being whether the *last* transaction in the sequence is fraud.
    """

    def __init__(self, sequences: np.ndarray, labels: np.ndarray):
        self.sequences = torch.FloatTensor(sequences)
        self.labels = torch.FloatTensor(labels)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.sequences[idx], self.labels[idx]


def build_sequences(
    df,
    feature_names: list,
    sequence_length: int = 20,
    cardholder_col: str = "cardholder_id",
    label_col: str = "is_fraud",
) -> tuple[np.ndarray, np.ndarray]:
    """Convert a transaction DataFrame into fixed-length sequences.

    For each transaction, we look back at the previous `sequence_length - 1`
    transactions by the same cardholder and form a sequence. Shorter histories
    are zero-padded on the left.

    Args:
        df: Transaction DataFrame sorted by timestamp.
        feature_names: Feature columns to include in sequences.
        sequence_length: Length of each transaction sequence.
        cardholder_col: Cardholder ID column name.
        label_col: Fraud label column name.

    Returns:
        Tuple of (sequences array [N, seq_len, n_features], labels array [N]).
    """
    df = df.sort_values([cardholder_col, "timestamp"]).reset_index(drop=True)
    n_features = len(feature_names)

    all_sequences = []
    all_labels = []

    for _, group in df.groupby(cardholder_col):
        features = group[feature_names].values
        labels = group[label_col].values

        for i in range(len(group)):
            start = max(0, i - sequence_length + 1)
            seq = features[start : i + 1]

            # Zero-pad if sequence is shorter than sequence_length
            if len(seq) < sequence_length:
                padding = np.zeros((sequence_length - len(seq), n_features))
                seq = np.vstack([padding, seq])

            all_sequences.append(seq)
            all_labels.append(labels[i])

    return np.array(all_sequences), np.array(all_labels)


# ---------------------------------------------------------------------------
# LSTM Model
# ---------------------------------------------------------------------------

class FraudLSTM(nn.Module):
    """Bidirectional LSTM for transaction sequence classification.

    Architecture:
    - Input embedding (linear projection)
    - 2-layer bidirectional LSTM
    - Attention pooling over sequence
    - Classification head with dropout
    """

    def __init__(
        self,
        input_dim: int,
        embedding_dim: int = 64,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.3,
    ):
        super().__init__()
        self.embedding = nn.Linear(input_dim, embedding_dim)
        self.lstm = nn.LSTM(
            input_size=embedding_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=True,
        )
        self.attention = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1),
        )
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            x: Input tensor of shape (batch, seq_len, input_dim).

        Returns:
            Fraud probability logits of shape (batch,).
        """
        embedded = self.embedding(x)  # (batch, seq_len, embed_dim)
        lstm_out, _ = self.lstm(embedded)  # (batch, seq_len, hidden*2)

        # Attention mechanism
        attn_weights = self.attention(lstm_out)  # (batch, seq_len, 1)
        attn_weights = torch.softmax(attn_weights, dim=1)
        context = (lstm_out * attn_weights).sum(dim=1)  # (batch, hidden*2)

        logits = self.classifier(context).squeeze(-1)  # (batch,)
        return logits


# ---------------------------------------------------------------------------
# Transformer Model
# ---------------------------------------------------------------------------

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for sequence position awareness."""

    def __init__(self, d_model: int, max_len: int = 500):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:, : x.size(1)]


class FraudTransformer(nn.Module):
    """Transformer encoder for transaction behavior profiling.

    Architecture:
    - Linear input projection + positional encoding
    - Multi-head self-attention encoder layers
    - CLS token pooling
    - Classification head
    """

    def __init__(
        self,
        input_dim: int,
        d_model: int = 128,
        nhead: int = 8,
        num_encoder_layers: int = 4,
        dim_feedforward: int = 256,
        dropout: float = 0.1,
        max_seq_len: int = 50,
    ):
        super().__init__()
        self.input_projection = nn.Linear(input_dim, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_seq_len)
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model))

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_encoder_layers
        )

        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model // 2, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            x: Input tensor of shape (batch, seq_len, input_dim).

        Returns:
            Fraud probability logits of shape (batch,).
        """
        batch_size = x.size(0)

        # Project input features to model dimension
        projected = self.input_projection(x)  # (batch, seq_len, d_model)
        projected = self.positional_encoding(projected)

        # Prepend CLS token
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        projected = torch.cat([cls_tokens, projected], dim=1)  # (batch, seq_len+1, d_model)

        # Transformer encoding
        encoded = self.transformer_encoder(projected)  # (batch, seq_len+1, d_model)

        # Use CLS token output for classification
        cls_output = encoded[:, 0]  # (batch, d_model)
        logits = self.classifier(cls_output).squeeze(-1)  # (batch,)
        return logits


# ---------------------------------------------------------------------------
# Training Loop
# ---------------------------------------------------------------------------

class DeepLearningTrainer:
    """Production training loop for deep learning fraud models.

    Features:
    - Mixed precision training
    - Learning rate scheduling with warmup
    - Early stopping on validation metric
    - Gradient clipping
    - Class-weighted loss for imbalanced data
    """

    def __init__(self, model: nn.Module, config: dict):
        self.model = model.to(DEVICE)
        self.config = config
        self.best_model_state = None
        self.training_history = []

    def train(
        self,
        train_sequences: np.ndarray,
        train_labels: np.ndarray,
        val_sequences: np.ndarray,
        val_labels: np.ndarray,
    ) -> dict:
        """Train the model with early stopping.

        Args:
            train_sequences: Training sequences (N, seq_len, features).
            train_labels: Training labels (N,).
            val_sequences: Validation sequences.
            val_labels: Validation labels.

        Returns:
            Dict with training history and best metrics.
        """
        batch_size = self.config.get("batch_size", 256)
        epochs = self.config.get("epochs", 50)
        lr = self.config.get("learning_rate", 0.001)
        patience = self.config.get("patience", 10)

        train_dataset = TransactionSequenceDataset(train_sequences, train_labels)
        val_dataset = TransactionSequenceDataset(val_sequences, val_labels)

        use_pin = DEVICE.type == "cuda"
        train_loader = DataLoader(
            train_dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=use_pin
        )
        val_loader = DataLoader(
            val_dataset, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=use_pin
        )

        # Class-weighted BCE loss (handle imbalance)
        pw_val = float((train_labels == 0).sum() / max(1, (train_labels == 1).sum()))
        pos_weight = torch.tensor([pw_val], device=DEVICE)
        criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=1e-5)

        # Cosine annealing with warmup
        warmup_steps = self.config.get("warmup_steps", 0)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

        # Metrics computed via sklearn on CPU numpy arrays (avoids PyTorch dispatch bugs)

        best_val_auc = 0.0
        patience_counter = 0

        logger.info("Starting training — %d epochs, batch_size=%d, lr=%s", epochs, batch_size, lr)
        logger.info("Device: %s, pos_weight: %.2f", DEVICE, pos_weight.item())

        for epoch in range(epochs):
            # --- Training ---
            self.model.train()
            train_loss = 0.0
            train_steps = 0

            for batch_x, batch_y in train_loader:
                batch_x = batch_x.to(DEVICE)
                batch_y = batch_y.to(DEVICE)

                optimizer.zero_grad()
                logits = self.model(batch_x)
                loss = criterion(logits, batch_y)
                loss.backward()

                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                optimizer.step()

                train_loss += loss.item()
                train_steps += 1

            scheduler.step()
            avg_train_loss = train_loss / max(1, train_steps)

            # --- Validation ---
            self.model.eval()
            val_loss = 0.0
            val_steps = 0
            all_logits = []
            all_labels = []

            with torch.no_grad():
                for batch_x, batch_y in val_loader:
                    batch_x = batch_x.to(DEVICE)
                    batch_y = batch_y.to(DEVICE)

                    logits = self.model(batch_x)
                    loss = criterion(logits, batch_y)
                    val_loss += loss.item()
                    val_steps += 1

                    all_logits.append(torch.sigmoid(logits))
                    all_labels.append(batch_y)

            avg_val_loss = val_loss / max(1, val_steps)
            all_probs_np = torch.cat(all_logits).cpu().numpy()
            all_labels_np = torch.cat(all_labels).cpu().numpy().astype(int)

            val_auc = roc_auc_score(all_labels_np, all_probs_np) if all_labels_np.sum() > 0 else 0.5
            val_ap = average_precision_score(all_labels_np, all_probs_np) if all_labels_np.sum() > 0 else 0.0

            self.training_history.append({
                "epoch": epoch + 1,
                "train_loss": avg_train_loss,
                "val_loss": avg_val_loss,
                "val_auc": val_auc,
                "val_ap": val_ap,
            })

            if (epoch + 1) % 5 == 0 or epoch == 0:
                logger.info(
                    "Epoch %d/%d — Train Loss: %.4f, Val Loss: %.4f, Val AUC: %.4f, Val AP: %.4f",
                    epoch + 1, epochs, avg_train_loss, avg_val_loss, val_auc, val_ap,
                )

            # Early stopping
            if val_auc > best_val_auc:
                best_val_auc = val_auc
                patience_counter = 0
                self.best_model_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    logger.info("Early stopping at epoch %d (best AUC: %.4f)", epoch + 1, best_val_auc)
                    break

        # Restore best model
        if self.best_model_state:
            self.model.load_state_dict(self.best_model_state)
            self.model.to(DEVICE)

        logger.info("Training complete — Best Val AUC: %.4f", best_val_auc)
        return {"best_val_auc": best_val_auc, "history": self.training_history}

    def predict_proba(self, sequences: np.ndarray) -> np.ndarray:
        """Predict fraud probabilities for transaction sequences.

        Args:
            sequences: Sequence array (N, seq_len, features).

        Returns:
            Fraud probability array (N,).
        """
        self.model.eval()
        dataset = TransactionSequenceDataset(sequences, np.zeros(len(sequences)))
        loader = DataLoader(dataset, batch_size=256, shuffle=False)

        all_probs = []
        with torch.no_grad():
            for batch_x, _ in loader:
                batch_x = batch_x.to(DEVICE)
                logits = self.model(batch_x)
                probs = torch.sigmoid(logits).cpu().numpy()
                all_probs.append(probs)

        return np.concatenate(all_probs)