"""
DM-BERT JEPA Architecture:
Integrates the bespoke 85.5M DM-BERT backbone with:
- Sub-quadratic Cartesian Block Butterfly Factorization (OpenAI 2026 Paper)
- First-class NULL_ACTION_ABSTAIN Candidate Head (Open-Set Abstention)
- Vectorized Candidate Option Encoder
- 4-step Latent Predictive Verifier in Thought Space
- Calibrated Cosine Compatibility Scorer
"""

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

from dm_bert.src.backbone import DMBertBackbone, BlockButterflyLinear, RMSNorm


class CriteriaEncoder(nn.Module):
    """
    Vectorized candidate criteria encoder.
    Encodes arbitrarily large candidate sets (K=2 to K=150) in < 0.2 ms
    using shared token embeddings and Block Butterfly linear projection.
    """
    def __init__(self, embed_tokens: nn.Embedding, d_model: int = 1024, p_butterfly: int = 32):
        super().__init__()
        self.embed_tokens = embed_tokens
        self.d_model = d_model
        self.proj = nn.Sequential(
            BlockButterflyLinear(d_model, d_model, p=p_butterfly),
            nn.GELU(),
            RMSNorm(d_model)
        )

    def forward(self, opt_input_ids: torch.Tensor, opt_attention_mask: torch.Tensor) -> torch.Tensor:
        B, K, L = opt_input_ids.shape
        flat_ids = opt_input_ids.view(B * K, L)
        flat_mask = opt_attention_mask.view(B * K, L)

        embeds = self.embed_tokens(flat_ids)  # [B*K, L, D]
        mask_exp = flat_mask.unsqueeze(-1).expand_as(embeds).float()
        sum_embeds = torch.sum(embeds * mask_exp, dim=1)
        denom = mask_exp.sum(dim=1).clamp(min=1e-9)
        pooled = sum_embeds / denom  # [B*K, D]

        out = self.proj(pooled)
        return out.view(B, K, self.d_model)


class LatentPredictorStep(nn.Module):
    """Single reasoning rollout step in thought space."""
    def __init__(self, d_model: int = 1024, nhead: int = 8, d_ff: int = 2048, dropout: float = 0.05):
        super().__init__()
        self.norm0 = RMSNorm(d_model)
        self.opt_self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.norm1 = RMSNorm(d_model)
        self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.norm2 = RMSNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        self.gate = nn.Sequential(
            nn.Linear(d_model * 2, d_model),
            nn.Sigmoid()
        )

    def forward(
        self,
        z: torch.Tensor,
        h_state: torch.Tensor,
        state_pad_mask: Optional[torch.Tensor] = None,
        opt_pad_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        # 1. Option interaction
        z_norm0 = self.norm0(z)
        attn0, _ = self.opt_self_attn(z_norm0, z_norm0, z_norm0, key_padding_mask=opt_pad_mask)
        z = z + attn0

        # 2. Cross attention to state facts
        z_norm1 = self.norm1(z)
        attn1, _ = self.cross_attn(z_norm1, h_state, h_state, key_padding_mask=state_pad_mask)
        z_attn = z + attn1

        # 3. Deduction
        z_norm2 = self.norm2(z_attn)
        ffn_out = self.ffn(z_norm2)

        # 4. Gated recurrent update
        cand = z_attn + ffn_out
        g = self.gate(torch.cat([z, cand], dim=-1))
        return g * cand + (1.0 - g) * z


class LatentPredictor(nn.Module):
    """M-step Latent Predictive Verifier with Block Butterfly output projection."""
    def __init__(self, d_model: int = 1024, nhead: int = 8, d_ff: int = 2048, num_steps: int = 4, p_butterfly: int = 32, dropout: float = 0.05):
        super().__init__()
        self.num_steps = num_steps
        self.cell = LatentPredictorStep(d_model, nhead, d_ff, dropout)
        self.out_proj = nn.Sequential(
            RMSNorm(d_model),
            BlockButterflyLinear(d_model, d_model, p=p_butterfly),
            RMSNorm(d_model)
        )

    def forward(
        self,
        h_state: torch.Tensor,
        s_options: torch.Tensor,
        state_attention_mask: Optional[torch.Tensor] = None,
        options_valid_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        state_pad_mask = (state_attention_mask == 0) if state_attention_mask is not None else None
        opt_pad_mask = (options_valid_mask == 0) if options_valid_mask is not None else None

        z = s_options
        for _ in range(self.num_steps):
            z = self.cell(z, h_state, state_pad_mask, opt_pad_mask)
        return self.out_proj(z)


class CompatibilityScorer(nn.Module):
    """Calibrated decision scoring via scaled cosine compatibility in thought space."""
    def __init__(self, d_model: int = 1024, init_temperature: float = 14.0):
        super().__init__()
        self.d_model = d_model
        self.log_temperature = nn.Parameter(torch.tensor(float(torch.log(torch.tensor(init_temperature)))))
        self.bias = nn.Parameter(torch.zeros(1))

    @property
    def temperature(self) -> torch.Tensor:
        return torch.exp(self.log_temperature)

    def forward(
        self,
        s_hat: torch.Tensor,
        s_options: torch.Tensor,
        options_valid_mask: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        norm_hat = F.normalize(s_hat, p=2, dim=-1)
        norm_opt = F.normalize(s_options, p=2, dim=-1)
        scores = torch.sum(norm_hat * norm_opt, dim=-1)
        logits = self.temperature * scores + self.bias
        if options_valid_mask is not None:
            logits = torch.where(
                options_valid_mask,
                logits,
                torch.tensor(-1e9, device=logits.device, dtype=logits.dtype)
            )
        probs = F.softmax(logits, dim=-1)
        return {"scores": scores, "logits": logits, "probabilities": probs}


class DMBertJEPA(nn.Module):
    """
    Complete DM-BERT JEPA Model:
    Bespoke 85.5M Bidirectional Transformer with Cartesian Block Butterfly layers
    and native NULL_ACTION_ABSTAIN open-set candidate head.
    """
    def __init__(
        self,
        vocab_size: int = 32768,
        d_model: int = 1024,
        n_layers: int = 12,
        n_heads: int = 16,
        num_steps: int = 4,
        p_butterfly: int = 32,
        init_temperature: float = 14.0
    ):
        super().__init__()
        self.d_model = d_model
        self.backbone = DMBertBackbone(
            vocab_size=vocab_size,
            d_model=d_model,
            n_layers=n_layers,
            n_heads=n_heads,
            p_butterfly=p_butterfly
        )
        self.criteria_encoder = CriteriaEncoder(self.backbone.embeddings, d_model=d_model, p_butterfly=p_butterfly)
        self.predictor = LatentPredictor(d_model=d_model, nhead=8, d_ff=2048, num_steps=num_steps, p_butterfly=p_butterfly)
        self.scorer = CompatibilityScorer(d_model=d_model, init_temperature=init_temperature)

        # First-class learnable NULL_ACTION_ABSTAIN candidate representation
        self.null_candidate = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

    def forward(
        self,
        state_input_ids: torch.Tensor,
        state_attention_mask: torch.Tensor,
        option_input_ids: torch.Tensor,
        option_attention_mask: torch.Tensor,
        options_valid_mask: Optional[torch.Tensor] = None,
        include_null: bool = False
    ) -> Dict[str, torch.Tensor]:
        # 1. State encoder
        h_state = self.backbone(state_input_ids, state_attention_mask)  # [B, L, 1024]

        # 2. Options encoder
        s_options = self.criteria_encoder(option_input_ids, option_attention_mask)  # [B, K, 1024]

        if include_null:
            B = s_options.size(0)
            null_exp = self.null_candidate.expand(B, 1, -1)
            s_options = torch.cat([s_options, null_exp], dim=1)
            if options_valid_mask is not None:
                null_valid = torch.ones((B, 1), dtype=torch.bool, device=options_valid_mask.device)
                options_valid_mask = torch.cat([options_valid_mask, null_valid], dim=1)

        # 3. Latent predictive rollout in thought space
        s_hat = self.predictor(h_state, s_options, state_attention_mask, options_valid_mask)

        # 4. Latent compatibility scoring
        out = self.scorer(s_hat, s_options, options_valid_mask)
        return out