"""Jue-Flash bidirectional encoder decision model with Multi-Layer Gating, Vectorized Attentive Span Pooling,
and Dual-Path (Global Context + Late-Interaction MaxSim) Readout for Choice / Score / Noul primitives.

Supports both ``v1`` checkpoints and ``v2`` architectural upgrades:
1. Query-Conditioned Multi-Depth Layer Gate across 6 transformer depths ``(-16, -12, -8, -4, -2, -1)``.
2. Multi-Hop Question/Candidate-to-State Cross-Attention refinement block.
3. Unimodal Ordinal Score Coupler combining continuous expected-severity regression with per-level span logits.
"""

from __future__ import annotations

import hashlib
import json
import math
from collections.abc import Mapping
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file, save_file
from torch import Tensor
from transformers import AutoConfig, AutoModel

from jue_flash.api.schema import MAX_SCORE_LEVELS
from jue_flash.encoding.pack import PackedFlashBatch

HEAD_WEIGHTS_FILE = "flash_decision_head.safetensors"
ENCODER_WEIGHTS_FILE = "encoder_model.safetensors"
FLASH_CONFIG_FILE = "jue_flash_config.json"


@dataclass
class FlashModelConfig:
    """Configuration for :class:`FlashDecisionModel`."""

    base_model: str = "jhu-clsp/mmbert-base"
    hidden_size: int = 768
    num_gated_layers: int = 4
    late_interaction_dim: int = 128
    dropout: float = 0.1
    max_len: int = 2048
    # v2 architectural flags (default False/0 so v1 checkpoints load identically):
    query_conditioned_layer_gate: bool = False
    tapped_layer_offsets: tuple[int, ...] = (-16, -12, -8, -4, -2, -1)
    cross_attn_hops: int = 0
    unimodal_score_coupler: bool = False
    # v3 architectural flags:
    canon_synergy_fusion: bool = False
    candidate_self_attn: bool = False

    def save(self, path: str | Path) -> None:
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        d = asdict(self)
        d["tapped_layer_offsets"] = list(self.tapped_layer_offsets)
        p.write_text(json.dumps(d, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    @classmethod
    def load(cls, path: str | Path) -> FlashModelConfig:
        raw = json.loads(Path(path).read_text(encoding="utf-8"))
        if "tapped_layer_offsets" in raw and isinstance(raw["tapped_layer_offsets"], list):
            raw["tapped_layer_offsets"] = tuple(int(x) for x in raw["tapped_layer_offsets"])
        return cls(**raw)


class LayerGate(nn.Module):
    """Learned softmax gate over the top ``num_layers`` hidden states of the encoder (v1)."""

    def __init__(self, num_layers: int, hidden_size: int) -> None:
        super().__init__()
        self.num_layers = num_layers
        self.layer_weights = nn.Parameter(torch.zeros(num_layers))
        self.norm = nn.LayerNorm(hidden_size)

    def forward(self, hidden_states: tuple[Tensor, ...]) -> Tensor:
        layers = hidden_states[-self.num_layers :]
        stacked = torch.stack(layers, dim=0)  # [L_g, B, T, D]
        w = torch.softmax(self.layer_weights, dim=0).view(-1, 1, 1, 1)
        fused = (stacked * w).sum(dim=0)
        return self.norm(fused)


class QueryConditionedLayerGate(nn.Module):
    """Dynamic per-request softmax gate across multi-depth encoder layers (v2).

    Taps shallow/middle syntactic layers and deep semantic layers (e.g. ``(-16, -12, -8, -4, -2, -1)``),
    conditioning the layer mixture weights on each request's top-layer summary vector so code/regex/AST
    questions can attend to intermediate layers while policy/reasoning questions attend to top layers.
    """

    def __init__(self, offsets: tuple[int, ...], hidden_size: int) -> None:
        super().__init__()
        self.offsets = offsets
        n_taps = len(offsets)
        self.layer_weights = nn.Parameter(torch.zeros(n_taps))
        self.query_norm = nn.LayerNorm(hidden_size)
        self.query_proj = nn.Linear(hidden_size, n_taps)
        nn.init.zeros_(self.query_proj.weight)
        nn.init.zeros_(self.query_proj.bias)
        self.norm = nn.LayerNorm(hidden_size)

    def forward(self, hidden_states: tuple[Tensor, ...]) -> Tensor:
        n_total = len(hidden_states)
        selected = [hidden_states[max(-n_total, min(-1, off))] for off in self.offsets]
        stacked = torch.stack(selected, dim=0)  # [L_g, B, T, D]
        # Condition on the top layer's CLS + mean summary per batch row: [B, D]
        top_h = hidden_states[-1]
        row_summary = 0.5 * (top_h[:, 0, :] + top_h.mean(dim=1))
        dyn_logits = self.layer_weights.unsqueeze(0) + self.query_proj(self.query_norm(row_summary))  # [B, L_g]
        w = torch.softmax(dyn_logits, dim=-1).transpose(0, 1).unsqueeze(-1).unsqueeze(-1)  # [L_g, B, 1, 1]
        fused = (stacked * w).sum(dim=0)
        return self.norm(fused)


class StateCrossAttentionHop(nn.Module):
    """Lightweight Multi-Head Cross-Attention hop refining ``[q_vec, c_1, ..., c_K]`` over state tokens (v2)."""

    def __init__(self, hidden_size: int, num_heads: int = 4, dropout: float = 0.1) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.q_norm = nn.LayerNorm(hidden_size)
        self.kv_norm = nn.LayerNorm(hidden_size)
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.out_proj = nn.Linear(hidden_size, hidden_size)
        nn.init.zeros_(self.out_proj.weight)
        nn.init.zeros_(self.out_proj.bias)
        self.dropout = nn.Dropout(dropout)

        self.ffn_norm = nn.LayerNorm(hidden_size)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 2),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size * 2, hidden_size),
        )
        nn.init.zeros_(self.ffn[-1].weight)  # type: ignore[arg-type]
        nn.init.zeros_(self.ffn[-1].bias)  # type: ignore[arg-type]

    def forward(
        self,
        qc_tokens: Tensor,  # [N, 1+K, D]
        state_seq_b: Tensor,  # [B, L, D]
        state_mask_b: Tensor,  # [B, L] bool
        q_b_idx: Tensor,  # [N]
    ) -> Tensor:
        n_q, n_qc, d = qc_tokens.shape
        bsz, seq_len, _ = state_seq_b.shape
        h = self.num_heads
        dh = self.head_dim

        kv_in = self.kv_norm(state_seq_b)
        k_b = self.k_proj(kv_in).view(bsz, seq_len, h, dh).transpose(1, 2)  # [B, H, L, dh]
        v_b = self.v_proj(kv_in).view(bsz, seq_len, h, dh).transpose(1, 2)  # [B, H, L, dh]

        q_in = self.q_norm(qc_tokens)
        q = self.q_proj(q_in).view(n_q, n_qc, h, dh).transpose(1, 2)  # [N, H, 1+K, dh]

        k = k_b[q_b_idx]  # [N, H, L, dh]
        v = v_b[q_b_idx]  # [N, H, L, dh]
        s_mask = state_mask_b[q_b_idx].unsqueeze(1).unsqueeze(2)  # [N, 1, 1, L]

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(dh)  # [N, H, 1+K, L]
        scores = scores.float().masked_fill(~s_mask, -1e9)
        attn = torch.softmax(scores, dim=-1).to(qc_tokens.dtype)
        attn = self.dropout(attn)

        context = torch.matmul(attn, v).transpose(1, 2).contiguous().view(n_q, n_qc, d)
        qc_tokens = qc_tokens + self.out_proj(context)
        qc_tokens = qc_tokens + self.ffn(self.ffn_norm(qc_tokens))
        return qc_tokens


class VectorizedAttentiveSpanPooler(nn.Module):
    """Attention + mean residual span pooler operating without Python loops over tokens."""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.attn_proj = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.Tanh(),
            nn.Linear(hidden_size // 2, 1, bias=False),
        )
        self.norm = nn.LayerNorm(hidden_size)

    def forward(self, token_states: Tensor, token_attn_scores: Tensor, spans: Tensor) -> Tensor:
        """Pool spans from ``token_states``."""
        m_size, seq_len, hidden_dim = token_states.shape
        pos = torch.arange(seq_len, device=token_states.device)

        if spans.dim() == 2:
            starts = spans[:, 0:1].clamp(0, seq_len - 1)
            ends = spans[:, 1:2].clamp(1, seq_len)
            ends = torch.maximum(ends, starts + 1)
            mask = (pos.unsqueeze(0) >= starts) & (pos.unsqueeze(0) < ends)  # [M, L]

            masked_scores = token_attn_scores.float().masked_fill(~mask, -1e9)
            weights = torch.softmax(masked_scores, dim=-1).to(token_states.dtype).unsqueeze(-1)  # [M, L, 1]
            attn_pooled = (token_states * weights).sum(dim=1)  # [M, D]

            mask_f = mask.to(token_states.dtype).unsqueeze(-1)
            mean_pooled = (token_states * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp_min(1.0)
            return self.norm(0.5 * (attn_pooled + mean_pooled))

        if spans.dim() == 3:
            starts = spans[:, :, 0:1].clamp(0, seq_len - 1)
            ends = spans[:, :, 1:2].clamp(1, seq_len)
            ends = torch.maximum(ends, starts + 1)
            p_grid = pos.view(1, 1, seq_len)
            mask = (p_grid >= starts) & (p_grid < ends)  # [M, K, L]

            scores_exp = token_attn_scores.float().unsqueeze(1).expand(-1, spans.shape[1], -1)  # [M, K, L]
            masked_scores = scores_exp.masked_fill(~mask, -1e9)
            weights = torch.softmax(masked_scores, dim=-1).to(token_states.dtype)  # [M, K, L]
            attn_pooled = torch.bmm(weights, token_states)  # [M, K, D]

            mask_f = mask.to(token_states.dtype)
            mean_weights = mask_f / mask_f.sum(dim=-1, keepdim=True).clamp_min(1.0)
            mean_pooled = torch.bmm(mean_weights, token_states)  # [M, K, D]
            return self.norm(0.5 * (attn_pooled + mean_pooled))

        raise ValueError(f"Expected spans with 2 or 3 dims, got {spans.shape}")


class CandidateSelfAttention(nn.Module):
    """Pure neural option-to-option self-attention and mean-contrast scorer for closed-book MCQs."""

    def __init__(self, hidden_size: int, num_heads: int = 4, dropout: float = 0.1) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.norm1 = nn.LayerNorm(hidden_size)
        self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=False)
        self.out_proj = nn.Linear(hidden_size, hidden_size)
        nn.init.zeros_(self.out_proj.weight)
        nn.init.zeros_(self.out_proj.bias)
        self.dropout = nn.Dropout(dropout)

        # Contrastive readout over [c_vec, c_diff, c_vec * q_exp, c_diff * q_exp] -> [N, K, 1]
        self.contrast_mlp = nn.Sequential(
            nn.Linear(4 * hidden_size, hidden_size // 2),
            nn.LayerNorm(hidden_size // 2),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, 1),
        )
        nn.init.zeros_(self.contrast_mlp[-1].weight)  # type: ignore[arg-type]
        nn.init.zeros_(self.contrast_mlp[-1].bias)  # type: ignore[arg-type]

        # Context-adaptive gate for late-interaction lexical similarity
        self.late_gate = nn.Sequential(
            nn.Linear(2 * hidden_size, hidden_size // 4),
            nn.LayerNorm(hidden_size // 4),
            nn.SiLU(),
            nn.Linear(hidden_size // 4, 1),
        )
        nn.init.zeros_(self.late_gate[-1].weight)  # type: ignore[arg-type]
        nn.init.zeros_(self.late_gate[-1].bias)  # type: ignore[arg-type]

    def forward(
        self,
        c_vec: Tensor,  # [N, K, D]
        q_vec: Tensor,  # [N, D]
        state_vec: Tensor,  # [N, D]
        cand_mask: Tensor,  # [N, K] bool
    ) -> tuple[Tensor, Tensor, Tensor]:
        n_q, k_max, d = c_vec.shape
        h = self.num_heads
        dh = self.head_dim

        x = self.norm1(c_vec)
        qkv = self.qkv(x).view(n_q, k_max, 3, h, dh).permute(2, 0, 3, 1, 4)  # [3, N, H, K, dh]
        q, k, v = qkv[0], qkv[1], qkv[2]
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(dh)  # [N, H, K, K]
        k_mask = cand_mask.unsqueeze(1).unsqueeze(2)  # [N, 1, 1, K]
        scores = scores.float().masked_fill(~k_mask, -1e9)
        attn = torch.softmax(scores, dim=-1).to(c_vec.dtype)
        attn = self.dropout(attn)
        ctx = torch.matmul(attn, v).transpose(1, 2).contiguous().view(n_q, k_max, d)
        c_refined = c_vec + self.out_proj(ctx)

        mask_f = cand_mask.to(c_vec.dtype).unsqueeze(-1)  # [N, K, 1]
        c_mean = (c_refined * mask_f).sum(dim=1, keepdim=True) / mask_f.sum(dim=1, keepdim=True).clamp_min(1.0)
        c_diff = c_refined - c_mean
        q_exp = q_vec.unsqueeze(1).expand(n_q, k_max, d)
        contrast_in = torch.cat([c_refined, c_diff, c_refined * q_exp, c_diff * q_exp], dim=-1)
        contrast_logits = self.contrast_mlp(contrast_in).squeeze(-1)  # [N, K]

        # Late-interaction gate in (0, 2), starting at 1.0 at initialization
        lg_in = torch.cat([q_vec, state_vec], dim=-1)
        late_mult = 1.0 + torch.tanh(self.late_gate(lg_in))  # [N, 1]
        return c_refined, contrast_logits, late_mult


class DualPathDecisionHead(nn.Module):
    """Three-primitive (Choice / Score / Noul) readout combining:
    - Path A: Global Contextual Interaction MLP + Bilinear scorer
    - Path B: Late-Interaction Token Alignment (ColBERT-style MaxSim between option tokens and state tokens)
    - Optional v2 upgrades: QueryConditionedLayerGate, StateCrossAttentionHop, and Unimodal Score Coupler.
    - Optional v3 upgrades: CandidateSelfAttention option cross-comparison.
    """

    def __init__(self, cfg: FlashModelConfig) -> None:
        super().__init__()
        self.cfg = cfg
        d = cfg.hidden_size
        if cfg.query_conditioned_layer_gate:
            self.layer_gate: nn.Module = QueryConditionedLayerGate(cfg.tapped_layer_offsets, d)
        else:
            self.layer_gate = LayerGate(cfg.num_gated_layers, d)
        self.span_pooler = VectorizedAttentiveSpanPooler(d)

        # Primitive embeddings: 0=choice, 1=score, 2=noul
        self.primitive_emb = nn.Embedding(3, d)
        # Ordinal level embeddings for Score (0..MAX_SCORE_LEVELS-1)
        self.score_level_emb = nn.Embedding(MAX_SCORE_LEVELS + 2, d)
        # Pole embeddings for Noul (0=false, 1=true)
        self.noul_pole_emb = nn.Embedding(2, d)

        # Optional v2 Multi-Hop Question/Candidate-to-State Cross-Attention
        if cfg.cross_attn_hops > 0:
            self.cross_hops = nn.ModuleList(
                [StateCrossAttentionHop(d, num_heads=4, dropout=cfg.dropout) for _ in range(cfg.cross_attn_hops)]
            )
        else:
            self.cross_hops = None

        # Optional v2 Unimodal Ordinal Score Coupler
        if cfg.unimodal_score_coupler:
            self.ordinal_regressor = nn.Sequential(
                nn.Linear(4 * d, d // 2),
                nn.LayerNorm(d // 2),
                nn.SiLU(),
                nn.Dropout(cfg.dropout),
                nn.Linear(d // 2, 1),
            )
            self.ordinal_log_precision = nn.Parameter(torch.tensor(math.log(3.0)))
        else:
            self.ordinal_regressor = None
            self.ordinal_log_precision = None

        # Path A: Global context-conditioned MLP + Bilinear interaction
        # Features: [c, q, s, c * q, c * s, |c - s|] -> 6 * D
        self.context_mlp = nn.Sequential(
            nn.Linear(6 * d, d),
            nn.LayerNorm(d),
            nn.SiLU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(d, d // 2),
            nn.SiLU(),
            nn.Linear(d // 2, 1),
        )
        self.bilinear_q = nn.Linear(d, d, bias=False)
        self.bilinear_s = nn.Linear(d, d, bias=False)

        # Path B: Late-Interaction token projection for exact identifier / telemetry / code alignment
        self.late_proj = nn.Linear(d, cfg.late_interaction_dim, bias=False)
        self.late_scale = nn.Parameter(torch.tensor(1.0))

        # Optional v3 Path C: Canon Specification Prior Synergy Gate (disabled in pure closed-book)
        self.canon_scale: nn.Parameter | None = None
        self.canon_gate: nn.Module | None = None
        if cfg.canon_synergy_fusion:
            self.enable_canon_synergy(init_scale=2.2)

        # Optional v3 Pure Neural Candidate Self-Attention & Contrast Block
        self.cand_attn: CandidateSelfAttention | None = None
        if cfg.candidate_self_attn:
            self.enable_candidate_self_attn()

    def enable_canon_synergy(self, init_scale: float = 2.2) -> None:
        """Initialize or re-enable the v3 Canon Specification Prior Synergy module."""
        d = self.cfg.hidden_size
        self.cfg.canon_synergy_fusion = True
        self.canon_scale = nn.Parameter(torch.tensor(float(init_scale)))
        gate = nn.Sequential(
            nn.Linear(2 * d + 2, d // 4),
            nn.LayerNorm(d // 4),
            nn.SiLU(),
            nn.Linear(d // 4, 1),
        )
        nn.init.zeros_(gate[-1].weight)  # type: ignore[arg-type]
        nn.init.zeros_(gate[-1].bias)  # type: ignore[arg-type]
        self.canon_gate = gate

    def enable_candidate_self_attn(self) -> None:
        """Initialize the pure neural CandidateSelfAttention option cross-comparison module."""
        self.cfg.candidate_self_attn = True
        self.cand_attn = CandidateSelfAttention(
            hidden_size=self.cfg.hidden_size,
            num_heads=4,
            dropout=self.cfg.dropout,
        )

    def _compute_late_interaction(
        self,
        late_tokens_b: Tensor,  # [B, L, d_late] normalized
        state_spans_b: Tensor,  # [B, 2]
        q_b_idx: Tensor,  # [N]
        opt_spans: Tensor,  # [N, K, 2]
        cand_mask: Tensor,  # [N, K]
    ) -> Tensor:
        """Compute ColBERT-style MaxSim between each candidate option's tokens and the state's tokens."""
        bsz, seq_len, _ = late_tokens_b.shape
        pos = torch.arange(seq_len, device=late_tokens_b.device)

        s_starts = state_spans_b[:, 0:1].clamp(0, seq_len - 1)
        s_ends = torch.maximum(state_spans_b[:, 1:2].clamp(1, seq_len), s_starts + 1)
        s_mask = (pos.unsqueeze(0) >= s_starts) & (pos.unsqueeze(0) < s_ends)  # [B, L]

        sim_matrix = torch.bmm(late_tokens_b, late_tokens_b.transpose(1, 2))  # [B, L_all, L_state]
        sim_matrix = sim_matrix.masked_fill(~s_mask.unsqueeze(1), -1e4)
        max_sim_to_state_b = sim_matrix.max(dim=-1).values  # [B, L]
        max_sim_to_state_q = max_sim_to_state_b[q_b_idx]  # [N, L]

        o_starts = opt_spans[:, :, 0:1].clamp(0, seq_len - 1)
        o_ends = torch.maximum(opt_spans[:, :, 1:2].clamp(1, seq_len), o_starts + 1)
        o_mask = (pos.view(1, 1, seq_len) >= o_starts) & (pos.view(1, 1, seq_len) < o_ends)  # [N, K, L]
        o_mask_f = o_mask.to(late_tokens_b.dtype)

        opt_maxsim = (o_mask_f * max_sim_to_state_q.unsqueeze(1)).sum(dim=-1) / o_mask_f.sum(dim=-1).clamp_min(1.0)
        return opt_maxsim.masked_fill(~cand_mask, 0.0) * self.late_scale

    def forward(self, hidden_states: tuple[Tensor, ...], batch: PackedFlashBatch) -> Tensor:
        """Compute candidate logits ``[N, K_max]`` for all questions in ``batch``."""
        fused_seq = self.layer_gate(hidden_states)  # [B, L, D]
        token_attn_scores = self.span_pooler.attn_proj(fused_seq).squeeze(-1)  # [B, L]

        # Pool state representation per batch row: [B, D]
        state_vec_b = self.span_pooler(fused_seq, token_attn_scores, batch.state_spans)  # [B, D]

        # Gather sequence representations for each question i in 0..N-1
        q_b_idx = batch.question_batch_idx  # [N]
        seq_for_q = fused_seq[q_b_idx]  # [N, L, D]
        scores_for_q = token_attn_scores[q_b_idx]  # [N, L]
        state_vec = state_vec_b[q_b_idx]  # [N, D]

        # Pool question instruction spans: [N, D]
        q_vec = self.span_pooler(seq_for_q, scores_for_q, batch.question_spans)  # [N, D]
        prim_vec = self.primitive_emb(batch.question_primitive_id)  # [N, D]
        q_vec = q_vec + prim_vec

        # Pool candidate option spans: [N, K, D]
        c_vec = self.span_pooler(seq_for_q, scores_for_q, batch.opt_spans)  # [N, K, D]
        n_q, k_max, d = c_vec.shape

        # Add ordinal level embeddings for Score questions and pole embeddings for Noul questions
        k_indices = torch.arange(k_max, device=c_vec.device).unsqueeze(0).expand(n_q, k_max)
        score_mask = (batch.question_primitive_id == 1).view(-1, 1, 1)  # [N, 1, 1]
        noul_mask = (batch.question_primitive_id == 2).view(-1, 1, 1)  # [N, 1, 1]

        lvl_emb = self.score_level_emb(k_indices.clamp(0, MAX_SCORE_LEVELS + 1))
        pole_emb = self.noul_pole_emb(k_indices.clamp(0, 1))
        c_vec = c_vec + score_mask.to(c_vec.dtype) * lvl_emb + noul_mask.to(c_vec.dtype) * pole_emb

        # Optional v2 Multi-Hop Cross-Attention over state tokens
        if self.cross_hops is not None:
            seq_len = fused_seq.shape[1]
            pos = torch.arange(seq_len, device=fused_seq.device)
            s_starts = batch.state_spans[:, 0:1].clamp(0, seq_len - 1)
            s_ends = torch.maximum(batch.state_spans[:, 1:2].clamp(1, seq_len), s_starts + 1)
            state_mask_b = (pos.unsqueeze(0) >= s_starts) & (pos.unsqueeze(0) < s_ends)  # [B, L]

            qc_tokens = torch.cat([q_vec.unsqueeze(1), c_vec], dim=1)  # [N, 1+K, D]
            for hop in self.cross_hops:
                qc_tokens = hop(qc_tokens, fused_seq, state_mask_b, q_b_idx)
            q_vec = qc_tokens[:, 0, :]
            c_vec = qc_tokens[:, 1:, :]

        # Optional v3 Pure Neural Candidate Self-Attention & Contrast
        contrast_logits: Tensor | None = None
        late_mult: Tensor | None = None
        if self.cand_attn is not None:
            c_vec, contrast_logits, late_mult = self.cand_attn(c_vec, q_vec, state_vec, batch.cand_mask)

        # Expand q_vec and state_vec to [N, K, D]
        q_exp = q_vec.unsqueeze(1).expand(n_q, k_max, d)
        s_exp = state_vec.unsqueeze(1).expand(n_q, k_max, d)

        # Path A: Global Contextual Interaction MLP + Bilinear
        features = torch.cat(
            [
                c_vec,
                q_exp,
                s_exp,
                c_vec * q_exp,
                c_vec * s_exp,
                (c_vec - s_exp).abs(),
            ],
            dim=-1,
        )  # [N, K, 6*D]
        mlp_logits = self.context_mlp(features).squeeze(-1)  # [N, K]
        bi_context = self.bilinear_q(q_vec) + self.bilinear_s(state_vec)  # [N, D]
        bi_logits = torch.bmm(c_vec, bi_context.unsqueeze(-1)).squeeze(-1) / (d**0.5)  # [N, K]

        # Path B: Late-Interaction MaxSim (computed once per state row B)
        late_tokens_b = F.normalize(self.late_proj(fused_seq), p=2, dim=-1)  # [B, L, d_late]
        late_logits = self._compute_late_interaction(
            late_tokens_b, batch.state_spans, q_b_idx, batch.opt_spans, batch.cand_mask
        )  # [N, K]
        if late_mult is not None:
            late_logits = late_logits * late_mult

        logits = mlp_logits + bi_logits + late_logits
        if contrast_logits is not None:
            logits = logits + contrast_logits

        # Optional v2 Unimodal Ordinal Score Coupler
        if self.ordinal_regressor is not None and self.ordinal_log_precision is not None:
            is_score = batch.question_primitive_id == 1  # [N]
            if is_score.any():
                ord_in = torch.cat([q_vec, state_vec, q_vec * state_vec, (q_vec - state_vec).abs()], dim=-1)  # [N, 4*D]
                pred_norm_lvl = torch.sigmoid(self.ordinal_regressor(ord_in))  # [N, 1] in (0, 1)
                num_valid_k = batch.cand_mask.sum(dim=-1, keepdim=True).clamp_min(2).to(logits.dtype) - 1.0  # [N, 1]
                lvl_coords = k_indices.to(logits.dtype) / num_valid_k  # [N, K] in [0, 1]
                prec = torch.exp(self.ordinal_log_precision.clamp(-2.0, 3.0))
                ord_bias = -prec * (lvl_coords - pred_norm_lvl).pow(2)  # [N, K]
                logits = logits + is_score.unsqueeze(-1).to(logits.dtype) * ord_bias

        return logits.masked_fill(~batch.cand_mask, float("-inf"))


class FlashDecisionModel(nn.Module):
    """Complete Jue-Flash model wrapping a Hugging Face bidirectional encoder and :class:`DualPathDecisionHead`."""

    def __init__(self, encoder: nn.Module, cfg: FlashModelConfig) -> None:
        super().__init__()
        self.encoder = encoder
        self.cfg = cfg
        self.head = DualPathDecisionHead(cfg)

    def upgrade_to_canon_synergy(self, init_scale: float = 2.2) -> None:
        """Enable v3 Canon Specification Prior Synergy Fusion on an initialized or loaded model."""
        device = next(self.head.parameters()).device
        dtype = next(self.head.parameters()).dtype
        self.head.enable_canon_synergy(init_scale=init_scale)
        self.head.to(device=device, dtype=dtype)

    def upgrade_to_candidate_self_attn(self) -> None:
        """Enable pure neural CandidateSelfAttention option cross-comparison on an initialized or loaded model."""
        device = next(self.head.parameters()).device
        dtype = next(self.head.parameters()).dtype
        self.head.enable_candidate_self_attn()
        self.head.to(device=device, dtype=dtype)

    @classmethod
    def from_pretrained_base(
        cls,
        base_model: str,
        *,
        num_gated_layers: int = 4,
        late_interaction_dim: int = 128,
        dropout: float = 0.1,
        max_len: int = 2048,
        head_v2: bool = True,
        canon_synergy_fusion: bool = False,
        candidate_self_attn: bool = False,
        dtype: torch.dtype = torch.float32,
        device: str | torch.device = "cpu",
        local_files_only: bool = False,
    ) -> FlashDecisionModel:
        """Initialize a fresh :class:`FlashDecisionModel` from a pretrained HF encoder backbone."""
        hf_cfg = AutoConfig.from_pretrained(base_model, local_files_only=local_files_only, trust_remote_code=True)
        hf_cfg.output_hidden_states = True
        if hasattr(hf_cfg, "reference_compile"):
            hf_cfg.reference_compile = False
        if getattr(hf_cfg, "model_type", "") == "lfm2" or "lfm2" in str(base_model).lower():
            from transformers import AutoModelForMaskedLM

            mlm = AutoModelForMaskedLM.from_pretrained(
                base_model,
                config=hf_cfg,
                dtype=dtype,
                local_files_only=local_files_only,
                trust_remote_code=True,
            )
            encoder = mlm.lfm2
        else:
            encoder = AutoModel.from_pretrained(
                base_model,
                config=hf_cfg,
                dtype=dtype,
                local_files_only=local_files_only,
                trust_remote_code=True,
            )
        cfg = FlashModelConfig(
            base_model=str(base_model),
            hidden_size=int(hf_cfg.hidden_size),
            num_gated_layers=min(num_gated_layers, int(getattr(hf_cfg, "num_hidden_layers", 12))),
            late_interaction_dim=late_interaction_dim,
            dropout=dropout,
            max_len=max_len,
            query_conditioned_layer_gate=head_v2,
            tapped_layer_offsets=(-16, -12, -8, -4, -2, -1),
            cross_attn_hops=1 if head_v2 else 0,
            unimodal_score_coupler=head_v2,
            canon_synergy_fusion=canon_synergy_fusion,
            candidate_self_attn=candidate_self_attn,
        )
        model = cls(encoder=encoder, cfg=cfg)
        return model.to(device=device, dtype=dtype)

    def save_checkpoint(self, out_dir: str | Path, *, encoder_dtype: torch.dtype | None = None) -> Path:
        """Save encoder weights, decision head weights, and FlashModelConfig to ``out_dir``."""
        d = Path(out_dir)
        d.mkdir(parents=True, exist_ok=True)
        self.cfg.save(d / FLASH_CONFIG_FILE)
        enc_state = {
            k: (v.detach().cpu().to(encoder_dtype) if encoder_dtype is not None and v.is_floating_point() else v.detach().cpu()).contiguous()
            for k, v in self.encoder.state_dict().items()
        }
        head_state = {k: v.detach().cpu().float().contiguous() for k, v in self.head.state_dict().items()}
        save_file(enc_state, str(d / ENCODER_WEIGHTS_FILE))
        save_file(head_state, str(d / HEAD_WEIGHTS_FILE))
        return d

    @classmethod
    def load_checkpoint(
        cls,
        checkpoint_dir: str | Path,
        *,
        base_model_override: str | None = None,
        dtype: torch.dtype = torch.float32,
        device: str | torch.device = "cpu",
        local_files_only: bool = False,
    ) -> FlashDecisionModel:
        """Load a saved :class:`FlashDecisionModel` checkpoint."""
        d = Path(checkpoint_dir)
        cfg = FlashModelConfig.load(d / FLASH_CONFIG_FILE)
        base_path = base_model_override or cfg.base_model
        hf_cfg = AutoConfig.from_pretrained(base_path, local_files_only=local_files_only, trust_remote_code=True)
        hf_cfg.output_hidden_states = True
        if hasattr(hf_cfg, "reference_compile"):
            hf_cfg.reference_compile = False
        encoder = AutoModel.from_config(hf_cfg, trust_remote_code=True)
        model = cls(encoder=encoder, cfg=cfg)
        enc_sd = load_file(str(d / ENCODER_WEIGHTS_FILE), device="cpu")
        head_sd = load_file(str(d / HEAD_WEIGHTS_FILE), device="cpu")
        model.encoder.load_state_dict(enc_sd, strict=True)
        model.head.load_state_dict(head_sd, strict=True)
        return model.to(device=device, dtype=dtype)

    def fingerprint(self) -> dict[str, Any]:
        """Compute a deterministic fingerprint of the model architecture and head parameters."""
        h = hashlib.sha256()
        for k, v in sorted(self.head.state_dict().items()):
            h.update(k.encode("utf-8"))
            h.update(v.detach().cpu().float().numpy().tobytes()[:256])
        return {
            "base_model": self.cfg.base_model,
            "hidden_size": self.cfg.hidden_size,
            "num_gated_layers": self.cfg.num_gated_layers,
            "late_interaction_dim": self.cfg.late_interaction_dim,
            "head_sha256_prefix": h.hexdigest()[:16],
        }

    def set_backbone_trainable(
        self,
        trainable: bool,
        unfreeze_top_layers: int | None = None,
        freeze_embeddings: bool = True,
    ) -> None:
        """Freeze/unfreeze encoder backbone parameters for Stage 1 head warmup vs Stage 2/3 full fine-tuning."""
        for p in self.encoder.parameters():
            p.requires_grad = trainable
        if freeze_embeddings:
            for attr in ("embeddings", "emb", "wte", "tok_embeddings", "embed_tokens"):
                emb_mod = getattr(self.encoder, attr, None)
                if emb_mod is not None:
                    for p in emb_mod.parameters():
                        p.requires_grad = False
        if not trainable and unfreeze_top_layers is not None and unfreeze_top_layers > 0:
            layers = getattr(self.encoder, "layers", None)
            if layers is None and hasattr(self.encoder, "encoder"):
                layers = getattr(self.encoder, "encoder", None)
                if layers is not None:
                    layers = getattr(layers, "layer", None)
            if layers is not None:
                for layer in list(layers)[-unfreeze_top_layers:]:
                    for p in layer.parameters():
                        p.requires_grad = True
        for p in self.head.parameters():
            p.requires_grad = True

    def forward(
        self,
        batch: PackedFlashBatch,
        *,
        temperature: float | Mapping[str, float] = 1.0,
        noul_bias: float = 0.0,
    ) -> Tensor:
        """Run single-pass bidirectional encoding and multi-question readout."""
        outputs = self.encoder(
            input_ids=batch.input_ids,
            attention_mask=batch.attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )
        logits = self.head(outputs.hidden_states, batch)  # [N, K_max]

        if noul_bias != 0.0 and logits.shape[1] >= 2:
            noul_rows = batch.question_primitive_id == 2
            if noul_rows.any():
                logits = logits.clone()
                logits[noul_rows, 1] = logits[noul_rows, 1] + float(noul_bias)

        if isinstance(temperature, Mapping):
            t_vec = torch.ones_like(batch.question_primitive_id, dtype=logits.dtype)
            t_vec = torch.where(
                batch.question_primitive_id == 0,
                torch.tensor(float(temperature.get("choice", 1.0)), dtype=logits.dtype, device=logits.device),
                t_vec,
            )
            t_vec = torch.where(
                batch.question_primitive_id == 1,
                torch.tensor(float(temperature.get("score", 1.0)), dtype=logits.dtype, device=logits.device),
                t_vec,
            )
            t_vec = torch.where(
                batch.question_primitive_id == 2,
                torch.tensor(float(temperature.get("noul", 1.0)), dtype=logits.dtype, device=logits.device),
                t_vec,
            )
            logits = logits / t_vec.unsqueeze(-1)
        elif temperature != 1.0:
            logits = logits / float(temperature)

        return logits.masked_fill(~batch.cand_mask, float("-inf"))