"""Dual-head ModernBERT for category + ordinal specificity classification.

Architecture:
    ModernBERT backbone → pooling (CLS or attention) → dropout →
        ├─ category head (7-class linear)
        └─ specificity head (options: coral, independent, softmax)

Specificity head options:
    - "coral": Single weight vector, 3 shifted biases (CoralLayer)
    - "independent": 3 independent binary classifiers, one per threshold
      (each learns its own feature direction — addresses CORAL limitation)
    - "softmax": Standard 4-class CE (no ordinal constraint)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from coral_pytorch.layers import CoralLayer


class FocalLoss(nn.Module):
    """Focal Loss (Lin et al. 2017) for imbalanced classification."""

    def __init__(self, weight: torch.Tensor | None = None, gamma: float = 2.0):
        super().__init__()
        self.gamma = gamma
        self.register_buffer("weight", weight)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        ce = F.cross_entropy(logits, targets, weight=self.weight, reduction="none")
        pt = torch.exp(-ce)
        return ((1 - pt) ** self.gamma) * ce


class AttentionPooling(nn.Module):
    """Learned attention pooling over all token representations."""

    def __init__(self, hidden_size: int):
        super().__init__()
        self.attention = nn.Linear(hidden_size, 1)

    def forward(
        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        scores = self.attention(hidden_states).squeeze(-1)  # (batch, seq_len)
        scores = scores.masked_fill(~attention_mask.bool(), float("-inf"))
        weights = F.softmax(scores, dim=1)
        return (hidden_states * weights.unsqueeze(-1)).sum(dim=1)


class IndependentThresholdHead(nn.Module):
    """Independent binary classifiers for each ordinal threshold.

    Unlike CORAL (shared weight vector, different biases), each threshold
    has its own fully independent Linear layer. This lets each transition
    learn different features:
        - L1→L2: cybersecurity terminology detection
        - L2→L3: firm-unique fact detection
        - L3→L4: quantified/verifiable claim detection
    """

    def __init__(self, in_features: int, num_classes: int, mlp_dim: int = 0):
        super().__init__()
        self.num_thresholds = num_classes - 1
        if mlp_dim > 0:
            self.thresholds = nn.ModuleList([
                nn.Sequential(
                    nn.Linear(in_features, mlp_dim),
                    nn.GELU(),
                    nn.Linear(mlp_dim, 1),
                )
                for _ in range(self.num_thresholds)
            ])
        else:
            self.thresholds = nn.ModuleList([
                nn.Linear(in_features, 1)
                for _ in range(self.num_thresholds)
            ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.cat([t(x) for t in self.thresholds], dim=1)


class DualHeadModernBERT(nn.Module):
    """ModernBERT with category classification + ordinal specificity heads."""

    def __init__(
        self,
        backbone: nn.Module,
        hidden_size: int,
        num_categories: int = 7,
        num_specificity: int = 4,
        dropout: float = 0.1,
        loss_type: str = "ce",
        focal_gamma: float = 2.0,
        category_weights: torch.Tensor | None = None,
        specificity_weights: list[float] | None = None,
        category_loss_weight: float = 1.0,
        specificity_loss_weight: float = 1.0,
        specificity_head_type: str = "independent",
        spec_mlp_dim: int = 0,
        pooling: str = "cls",
        ordinal_consistency_weight: float = 0.0,
    ):
        super().__init__()
        self.backbone = backbone
        self.dropout = nn.Dropout(dropout)
        self.num_specificity = num_specificity
        self.category_loss_weight = category_loss_weight
        self.specificity_loss_weight = specificity_loss_weight
        self.specificity_head_type = specificity_head_type
        self.ordinal_consistency_weight = ordinal_consistency_weight

        # Pooling
        self.pooling_type = pooling
        if pooling == "attention":
            self.pooler = AttentionPooling(hidden_size)

        # Category head
        self.category_head = nn.Linear(hidden_size, num_categories)

        # Specificity head
        if specificity_head_type == "coral":
            self.specificity_head = CoralLayer(hidden_size, num_specificity)
        elif specificity_head_type == "independent":
            self.specificity_head = IndependentThresholdHead(
                hidden_size, num_specificity, mlp_dim=spec_mlp_dim
            )
        elif specificity_head_type == "softmax":
            self.specificity_head = nn.Linear(hidden_size, num_specificity)
        else:
            raise ValueError(f"Unknown specificity_head_type: {specificity_head_type}")

        # Category loss
        if loss_type == "focal":
            self.category_loss_fn = FocalLoss(weight=category_weights, gamma=focal_gamma)
        else:
            self.category_loss_fn = nn.CrossEntropyLoss(weight=category_weights, reduction="none")

        self.specificity_weights = specificity_weights

    def _pool(
        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        if self.pooling_type == "attention":
            return self.pooler(hidden_states, attention_mask)
        return hidden_states[:, 0]  # [CLS]

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        category_labels: torch.Tensor | None = None,
        specificity_labels: torch.Tensor | None = None,
        sample_weight: torch.Tensor | None = None,
        spec_mask: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        pooled = self._pool(outputs.last_hidden_state, attention_mask)
        pooled = self.dropout(pooled)

        cat_logits = self.category_head(pooled)
        spec_logits = self.specificity_head(pooled)

        result = {
            "category_logits": cat_logits,
            "specificity_logits": spec_logits,
        }

        if category_labels is not None and specificity_labels is not None:
            # Category loss (all samples)
            cat_loss_per_sample = self.category_loss_fn(cat_logits, category_labels)
            if sample_weight is not None:
                cat_loss_per_sample = cat_loss_per_sample * sample_weight
            cat_loss = cat_loss_per_sample.mean()

            # Specificity loss
            if self.specificity_head_type == "softmax":
                spec_loss_per_sample = F.cross_entropy(
                    spec_logits, specificity_labels, reduction="none"
                )
            else:
                # Cumulative binary CE (works for both coral and independent)
                spec_loss_per_sample = _cumulative_bce_per_sample(
                    spec_logits, specificity_labels, self.num_specificity
                )

            if sample_weight is not None:
                spec_loss_per_sample = spec_loss_per_sample * sample_weight

            # Apply spec_mask: only compute spec loss on high-confidence samples
            if spec_mask is not None:
                spec_loss_per_sample = spec_loss_per_sample * spec_mask
                spec_loss = spec_loss_per_sample.sum() / (spec_mask.sum() + 1e-8)
            else:
                spec_loss = spec_loss_per_sample.mean()

            total_loss = (
                self.category_loss_weight * cat_loss
                + self.specificity_loss_weight * spec_loss
            )

            # Ordinal consistency regularization (for independent thresholds)
            if self.ordinal_consistency_weight > 0 and self.specificity_head_type != "softmax":
                consistency_loss = _ordinal_consistency_loss(spec_logits)
                total_loss = total_loss + self.ordinal_consistency_weight * consistency_loss

            result["loss"] = total_loss
            result["category_loss"] = cat_loss.detach()
            result["specificity_loss"] = spec_loss.detach()

        return result


def _cumulative_bce_per_sample(
    logits: torch.Tensor, labels: torch.Tensor, num_classes: int
) -> torch.Tensor:
    """Per-sample cumulative binary CE (shared by CORAL and independent heads)."""
    levels = torch.zeros(
        labels.size(0), num_classes - 1, device=logits.device, dtype=logits.dtype
    )
    for k in range(num_classes - 1):
        levels[:, k] = (labels > k).float()
    return F.binary_cross_entropy_with_logits(
        logits, levels, reduction="none"
    ).sum(dim=1)


def _ordinal_consistency_loss(logits: torch.Tensor) -> torch.Tensor:
    """Penalize threshold k firing when threshold k-1 doesn't.

    For cumulative thresholds, P(L≥3) should never exceed P(L≥2).
    Penalty: max(0, sigmoid(logit_k) - sigmoid(logit_{k-1})) for all k.
    """
    probs = torch.sigmoid(logits)  # (batch, num_thresholds)
    # probs[:, k] should be <= probs[:, k-1]
    violations = F.relu(probs[:, 1:] - probs[:, :-1])
    return violations.mean()


def ordinal_predict(logits: torch.Tensor) -> torch.Tensor:
    """Convert cumulative threshold logits to class predictions (0-indexed)."""
    probs = torch.sigmoid(logits)
    return (probs > 0.5).sum(dim=1).long()


def softmax_predict(logits: torch.Tensor) -> torch.Tensor:
    """Convert softmax logits to class predictions."""
    return logits.argmax(dim=1)
