"""
model.py – TRALSem: Textual Robust Adversarial Learning for Sentiment Analysis
=================================================================================
Architecture (Paper §3):
    Input
      → BERT Tokeniser (token_ids, segment_ids, position_ids)
      → Data Encoder   (token + position + segment embeddings)
      → Transformer Encoder (12 layers, multi-head self-attention)
      → [CLS] pooling
      → Linear Layer   (regression-style projection)
      → Softmax        (classification output)

The Random-Seed mechanism (LCG) ensures reproducibility, and adversarial
perturbations are injected at the embedding level during training (see
adversarial.py and train.py).
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from config import (
    HIDDEN_SIZE, NUM_ATTENTION_HEADS, NUM_HIDDEN_LAYERS,
    INTERMEDIATE_SIZE, HIDDEN_DROPOUT, ATTENTION_DROPOUT,
    MAX_SEQ_LEN, LCG_A, LCG_C, LCG_M, LCG_SEED,
)


# ═══════════════════════════════════════════════════════════════════════════════
# Random Seed Generator (LCG) – Paper §3.2
# ═══════════════════════════════════════════════════════════════════════════════
class LCGRandomSeed:
    """
    Linear Congruential Generator for deterministic seed production.
        s_{n+1} = (a · s_n + c) mod m
    Used to set random seeds at each training step for exact reproducibility.
    """

    def __init__(self, a: int = LCG_A, c: int = LCG_C,
                 m: int = LCG_M, seed: int = LCG_SEED):
        self.a = a
        self.c = c
        self.m = m
        self.state = seed

    def next(self) -> int:
        self.state = (self.a * self.state + self.c) % self.m
        return self.state

    def set_global_seed(self):
        """Push deterministic seed into PyTorch + Python."""
        s = self.next()
        torch.manual_seed(s)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(s)
        return s


# ═══════════════════════════════════════════════════════════════════════════════
# Multi-Head Self-Attention – Paper Eq: Attention(Q,K,V)=softmax(QK^T/√d_k)V
# ═══════════════════════════════════════════════════════════════════════════════
class MultiHeadSelfAttention(nn.Module):
    """Exact implementation of scaled dot-product multi-head attention."""

    def __init__(self, hidden_size: int = HIDDEN_SIZE,
                 num_heads: int = NUM_ATTENTION_HEADS,
                 dropout: float = ATTENTION_DROPOUT):
        super().__init__()
        assert hidden_size % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads  # d_k

        self.query = nn.Linear(hidden_size, hidden_size)
        self.key   = nn.Linear(hidden_size, hidden_size)
        self.value = nn.Linear(hidden_size, hidden_size)
        self.out   = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, hidden_states: torch.Tensor,
                attention_mask: torch.Tensor = None) -> torch.Tensor:
        B, L, _ = hidden_states.size()

        # Linear projections → (B, num_heads, L, head_dim)
        Q = self.query(hidden_states).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.key(hidden_states).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.value(hidden_states).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)

        if attention_mask is not None:
            # attention_mask: (B, 1, 1, L) — broadcast across heads & queries
            scores = scores + attention_mask

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        context = torch.matmul(attn_weights, V)  # (B, H, L, d_k)
        context = context.transpose(1, 2).contiguous().view(B, L, -1)
        return self.out(context)


# ═══════════════════════════════════════════════════════════════════════════════
# Transformer Encoder Block (one layer)
# ═══════════════════════════════════════════════════════════════════════════════
class TransformerBlock(nn.Module):
    """Pre-LN transformer block: Attention → Add & Norm → FFN → Add & Norm."""

    def __init__(self, hidden_size: int = HIDDEN_SIZE,
                 num_heads: int = NUM_ATTENTION_HEADS,
                 intermediate_size: int = INTERMEDIATE_SIZE,
                 dropout: float = HIDDEN_DROPOUT):
        super().__init__()
        self.attention = MultiHeadSelfAttention(hidden_size, num_heads)
        self.norm1 = nn.LayerNorm(hidden_size)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.GELU(),
            nn.Linear(intermediate_size, hidden_size),
        )
        self.norm2 = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, attention_mask=None) -> torch.Tensor:
        # Self-attention sub-layer
        attn_out = self.attention(self.norm1(x), attention_mask)
        x = x + self.dropout(attn_out)
        # Feed-forward sub-layer
        ffn_out = self.ffn(self.norm2(x))
        x = x + self.dropout(ffn_out)
        return x


# ═══════════════════════════════════════════════════════════════════════════════
# Data Encoder (Token + Position + Segment embeddings) – Paper §3.1
# ═══════════════════════════════════════════════════════════════════════════════
class DataEncoder(nn.Module):
    """
    Converts token ids, position ids, and segment ids into a continuous
    embedding E = E_token + E_position + E_segment, followed by LayerNorm
    and dropout — identical to the BERT embedding layer.
    """

    def __init__(self, vocab_size: int = 30522,
                 hidden_size: int = HIDDEN_SIZE,
                 max_position: int = MAX_SEQ_LEN,
                 type_vocab_size: int = 2,
                 dropout: float = HIDDEN_DROPOUT):
        super().__init__()
        self.word_embeddings     = nn.Embedding(vocab_size, hidden_size)
        self.position_embeddings = nn.Embedding(max_position, hidden_size)
        self.token_type_embeddings = nn.Embedding(type_vocab_size, hidden_size)
        self.norm = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input_ids: torch.Tensor,
                token_type_ids: torch.Tensor = None) -> torch.Tensor:
        B, L = input_ids.size()
        position_ids = torch.arange(L, device=input_ids.device).unsqueeze(0).expand(B, -1)
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        embeddings = (
            self.word_embeddings(input_ids)
            + self.position_embeddings(position_ids)
            + self.token_type_embeddings(token_type_ids)
        )
        return self.dropout(self.norm(embeddings))


# ═══════════════════════════════════════════════════════════════════════════════
# TRALSem – Full Model
# ═══════════════════════════════════════════════════════════════════════════════
class TRALSem(nn.Module):
    """
    TRALSem: Textual Robust Adversarial Learning for Sentiment Analysis.

    Pipeline:
        Input → DataEncoder → [Transformer × 12] → [CLS] → LinearLayer → Softmax

    The Linear Layer produces a regression-style scalar per class (paper §3.4),
    and Softmax converts to probabilities.  RMSE loss is used during training.
    """

    def __init__(self, vocab_size: int = 30522,
                 num_labels: int = 5,
                 hidden_size: int = HIDDEN_SIZE,
                 num_layers: int = NUM_HIDDEN_LAYERS,
                 num_heads: int = NUM_ATTENTION_HEADS,
                 intermediate_size: int = INTERMEDIATE_SIZE,
                 dropout: float = HIDDEN_DROPOUT,
                 max_seq_len: int = MAX_SEQ_LEN):
        super().__init__()
        self.num_labels = num_labels

        # ── Data Encoder ─────────────────────────────────────────────────
        self.encoder = DataEncoder(vocab_size, hidden_size, max_seq_len,
                                   dropout=dropout)

        # ── Transformer Encoder (12 layers) ──────────────────────────────
        self.transformer_layers = nn.ModuleList([
            TransformerBlock(hidden_size, num_heads, intermediate_size, dropout)
            for _ in range(num_layers)
        ])
        self.final_norm = nn.LayerNorm(hidden_size)

        # ── Linear Layer (LL) – regression-style head (Paper §3.4) ──────
        self.classifier = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Tanh(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, num_labels),
        )

        self._init_weights()

    # ── weight initialisation ─────────────────────────────────────────────
    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

    # ── load pre-trained BERT weights ─────────────────────────────────────
    def load_pretrained_bert(self, bert_model):
        """
        Copy weights from a HuggingFace BertModel into our custom layers.
        This bridges the gap between our from-scratch architecture and
        pretrained BERT for transfer learning.
        """
        # Embeddings
        self.encoder.word_embeddings.weight.data.copy_(
            bert_model.embeddings.word_embeddings.weight.data)
        self.encoder.position_embeddings.weight.data.copy_(
            bert_model.embeddings.position_embeddings.weight.data[:MAX_SEQ_LEN])
        self.encoder.token_type_embeddings.weight.data.copy_(
            bert_model.embeddings.token_type_embeddings.weight.data)
        self.encoder.norm.weight.data.copy_(
            bert_model.embeddings.LayerNorm.weight.data)
        self.encoder.norm.bias.data.copy_(
            bert_model.embeddings.LayerNorm.bias.data)

        # Transformer layers
        for i, layer in enumerate(self.transformer_layers):
            bl = bert_model.encoder.layer[i]
            # Attention
            layer.attention.query.weight.data.copy_(bl.attention.self.query.weight.data)
            layer.attention.query.bias.data.copy_(bl.attention.self.query.bias.data)
            layer.attention.key.weight.data.copy_(bl.attention.self.key.weight.data)
            layer.attention.key.bias.data.copy_(bl.attention.self.key.bias.data)
            layer.attention.value.weight.data.copy_(bl.attention.self.value.weight.data)
            layer.attention.value.bias.data.copy_(bl.attention.self.value.bias.data)
            layer.attention.out.weight.data.copy_(bl.attention.output.dense.weight.data)
            layer.attention.out.bias.data.copy_(bl.attention.output.dense.bias.data)
            # LayerNorm 1
            layer.norm1.weight.data.copy_(bl.attention.output.LayerNorm.weight.data)
            layer.norm1.bias.data.copy_(bl.attention.output.LayerNorm.bias.data)
            # FFN
            layer.ffn[0].weight.data.copy_(bl.intermediate.dense.weight.data)
            layer.ffn[0].bias.data.copy_(bl.intermediate.dense.bias.data)
            layer.ffn[2].weight.data.copy_(bl.output.dense.weight.data)
            layer.ffn[2].bias.data.copy_(bl.output.dense.bias.data)
            # LayerNorm 2
            layer.norm2.weight.data.copy_(bl.output.LayerNorm.weight.data)
            layer.norm2.bias.data.copy_(bl.output.LayerNorm.bias.data)

    # ── forward pass ──────────────────────────────────────────────────────
    def forward(self, input_ids: torch.Tensor,
                attention_mask: torch.Tensor = None,
                token_type_ids: torch.Tensor = None) -> torch.Tensor:
        """
        Returns raw logits (before softmax) of shape (B, num_labels).
        Softmax is applied externally in loss / inference.
        """
        # 1. Data Encoder
        hidden = self.encoder(input_ids, token_type_ids)

        # 2. Prepare attention mask → (B, 1, 1, L), -10000 for padded tokens
        if attention_mask is not None:
            extended_mask = attention_mask[:, None, None, :]
            extended_mask = (1.0 - extended_mask) * -1e4
        else:
            extended_mask = None

        # 3. Transformer Encoder (12 layers)
        for layer in self.transformer_layers:
            hidden = layer(hidden, extended_mask)
        hidden = self.final_norm(hidden)

        # 4. [CLS] token pooling
        cls_output = hidden[:, 0]

        # 5. Linear Layer → logits
        logits = self.classifier(cls_output)

        return logits

    # ── convenience: softmax probabilities ────────────────────────────────
    def predict_proba(self, input_ids, attention_mask=None,
                      token_type_ids=None) -> torch.Tensor:
        logits = self.forward(input_ids, attention_mask, token_type_ids)
        return F.softmax(logits, dim=-1)


# ═══════════════════════════════════════════════════════════════════════════════
# RMSE Loss – Paper §3.5
# ═══════════════════════════════════════════════════════════════════════════════
class RMSELoss(nn.Module):
    """
    Root Mean Squared Error loss operating on softmax probabilities.
    Converts labels to one-hot and computes RMSE between predicted
    probability vector and target vector.
    """

    def __init__(self, num_classes: int):
        super().__init__()
        self.num_classes = num_classes

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        probs = F.softmax(logits, dim=-1)
        one_hot = F.one_hot(targets, num_classes=self.num_classes).float()
        mse = torch.mean((probs - one_hot) ** 2)
        return torch.sqrt(mse + 1e-8)


# ═══════════════════════════════════════════════════════════════════════════════
# Forward-function helper for adversarial modules
# ═══════════════════════════════════════════════════════════════════════════════
def model_forward(model, input_ids, attention_mask, token_type_ids):
    """Plain forward — used as a callable by FreeLB / YOPO."""
    return model(input_ids, attention_mask, token_type_ids)