"""
Module B: Context-aware Risk Detector.

Architecture:
  1. Encode persona, context (history+user_input), response separately
  2. Fuse via CrossAttention(response, [persona; context])
  3. Multi-task classification heads:
     - Binary risk (sigmoid)
     - Risk level 0-4 (softmax)
     - Primary category R1-R10 (softmax)
     - Fine-grained 14-label (sigmoid multi-label)

Returns e_P_pool and e_H_pool for downstream RL state construction.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional

from src.models.encoder import TextEncoder, ContextAwareFusion
from src.utils.taxonomy import NUM_PRIMARY, NUM_FINE, NUM_RISK_LEVELS


class CompanionRiskDetector(nn.Module):
    def __init__(
        self,
        model_name: str = "hfl/chinese-macbert-large",
        hidden_size: int = 768,
        num_heads: int = 8,
        dropout: float = 0.1,
        use_lora: bool = False,
    ):
        super().__init__()

        self.encoder = TextEncoder(
            model_name=model_name,
            hidden_size=hidden_size,
            use_lora=use_lora,
        )

        self.fusion = ContextAwareFusion(
            hidden_size=hidden_size,
            num_heads=num_heads,
            dropout=dropout,
        )

        self.dropout = nn.Dropout(dropout)

        # Classification heads
        self.binary_head = nn.Linear(hidden_size, 1)
        self.level_head = nn.Linear(hidden_size, NUM_RISK_LEVELS)
        self.primary_head = nn.Linear(hidden_size, NUM_PRIMARY)
        self.fine_head = nn.Linear(hidden_size, NUM_FINE)

    def _mean_pool(self, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """Mean-pool token representations using attention mask."""
        m = mask.unsqueeze(-1).float()
        return (hidden * m).sum(1) / m.sum(1).clamp(min=1e-9)

    def _build_context_padding_mask(
        self,
        persona_mask: torch.Tensor,
        context_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Build boolean padding mask for CrossAttention (True = ignore position)."""
        return torch.cat([persona_mask == 0, context_mask == 0], dim=1)

    def forward(
        self,
        persona_input_ids: torch.Tensor,
        persona_attention_mask: torch.Tensor,
        context_input_ids: torch.Tensor,
        context_attention_mask: torch.Tensor,
        response_input_ids: torch.Tensor,
        response_attention_mask: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:

        # Encode all three streams — [B, seq_len, H]
        persona_h = self.encoder(persona_input_ids, persona_attention_mask)
        context_h = self.encoder(context_input_ids, context_attention_mask)
        response_h = self.encoder(response_input_ids, response_attention_mask)

        # Separate pooled representations for RL state
        e_P_pool = self._mean_pool(persona_h, persona_attention_mask)   # [B, H]
        e_H_pool = self._mean_pool(context_h, context_attention_mask)   # [B, H]

        # CrossAttention: response queries [persona; context] as relational context
        combined_context = torch.cat([persona_h, context_h], dim=1)
        combined_mask = self._build_context_padding_mask(persona_attention_mask, context_attention_mask)
        fused = self.fusion(response_h, combined_context, combined_mask)

        # Pool fused representation
        e_fused = self._mean_pool(fused, response_attention_mask)
        e_fused = self.dropout(e_fused)

        return {
            "y_risk": self.binary_head(e_fused).squeeze(-1),    # [B]
            "l_risk": self.level_head(e_fused),                  # [B, 5]
            "c_primary": self.primary_head(e_fused),             # [B, 10]
            "c_fine": self.fine_head(e_fused),                   # [B, 14]
            "e_fused": e_fused,                                  # [B, H]
            "e_P_pool": e_P_pool,                                # [B, H]
            "e_H_pool": e_H_pool,                                # [B, H]
        }

    def compute_loss(
        self,
        logits: Dict[str, torch.Tensor],
        targets: Dict[str, torch.Tensor],
        weights: Dict[str, float] = None,
        fine_pos_weight: Optional[torch.Tensor] = None,
        fine_risky_only: bool = False,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Args:
            fine_pos_weight: shape [NUM_FINE], class-specific positive weights for BCEWithLogitsLoss.
                             Computed from training set: pos_weight[i] = (N - pos_i) / pos_i.
                             Strongly recommended for next training round to handle rare labels
                             like Romanticization/CoRumination (pos_weight ≈ 25.8).
            fine_risky_only: if True, compute fine-label loss only on y_risk=1 samples.
                             Safe samples have empty c_fine by design; including them
                             in the loss teaches the model to predict all-negative, causing
                             fine_macro_f1 ≈ 0 at evaluation time.
        """
        if weights is None:
            weights = {"binary": 1.0, "level": 1.0, "primary": 1.0, "fine": 1.0}

        loss_parts = {}

        loss_binary = F.binary_cross_entropy_with_logits(
            logits["y_risk"], targets["y_risk"].float()
        )
        loss_parts["loss_binary"] = loss_binary

        loss_level = F.cross_entropy(logits["l_risk"], targets["l_risk"].long())
        loss_parts["loss_level"] = loss_level

        # Only compute primary category loss for samples with a valid category
        # c_primary target is one-hot; samples with c_primary = "None" have all-zero vectors
        primary_valid_mask = targets["c_primary"].sum(-1) > 0  # [B]
        if primary_valid_mask.any():
            primary_targets = targets["c_primary"][primary_valid_mask].argmax(-1)
            primary_logits = logits["c_primary"][primary_valid_mask]
            loss_primary = F.cross_entropy(primary_logits, primary_targets)
        else:
            loss_primary = torch.tensor(0.0, device=logits["c_primary"].device)
        loss_parts["loss_primary"] = loss_primary

        # Fine-grained multi-label loss
        # Optional: restrict to risky samples to avoid teaching all-negative on safe samples
        if fine_risky_only:
            risky_mask = targets["y_risk"] > 0.5  # [B]
            if risky_mask.any():
                fine_logits_masked = logits["c_fine"][risky_mask]
                fine_targets_masked = targets["c_fine"][risky_mask].float()
                loss_fine = F.binary_cross_entropy_with_logits(
                    fine_logits_masked, fine_targets_masked,
                    pos_weight=fine_pos_weight,
                )
            else:
                loss_fine = torch.tensor(0.0, device=logits["c_fine"].device)
        else:
            loss_fine = F.binary_cross_entropy_with_logits(
                logits["c_fine"], targets["c_fine"].float(),
                pos_weight=fine_pos_weight,
            )
        loss_parts["loss_fine"] = loss_fine

        total = (
            weights["binary"] * loss_binary
            + weights["level"] * loss_level
            + weights["primary"] * loss_primary
            + weights["fine"] * loss_fine
        )

        return total, loss_parts

    @torch.no_grad()
    def predict(
        self,
        persona_input_ids: torch.Tensor,
        persona_attention_mask: torch.Tensor,
        context_input_ids: torch.Tensor,
        context_attention_mask: torch.Tensor,
        response_input_ids: torch.Tensor,
        response_attention_mask: torch.Tensor,
        binary_threshold: float = 0.5,
        fine_threshold: float = 0.4,
    ) -> Dict:
        logits = self.forward(
            persona_input_ids, persona_attention_mask,
            context_input_ids, context_attention_mask,
            response_input_ids, response_attention_mask,
        )

        d_score = torch.sigmoid(logits["y_risk"])
        y_risk = (d_score >= binary_threshold).long()
        l_risk = logits["l_risk"].argmax(-1)
        c_primary = logits["c_primary"].argmax(-1)
        c_primary_probs = torch.softmax(logits["c_primary"], dim=-1)
        c_fine_probs = torch.sigmoid(logits["c_fine"])           # [B, NUM_FINE] continuous scores
        c_fine = (c_fine_probs >= fine_threshold).float()        # [B, NUM_FINE] binary predictions

        return {
            "y_risk": y_risk,
            "l_risk": l_risk,
            "c_primary": c_primary,
            "c_primary_probs": c_primary_probs,
            # c_fine: binary predictions (already thresholded) — use this in evaluate.py
            "c_fine": c_fine,
            # c_fine_probs: continuous sigmoid scores — use for ranking or custom thresholds
            "c_fine_probs": c_fine_probs,
            "d_score": d_score,
            "e_fused": logits["e_fused"],
            "e_P_pool": logits["e_P_pool"],
            "e_H_pool": logits["e_H_pool"],
        }
