"""MiniJEV Core: the domain-agnostic decision network.

It only sees vectors in a shared latent space of width `d`:

    state_tokens    (B, S, d) + state_mask (B, S)      from any StateEncoder (grid, text, numeric, ...)
    question_tokens (B, Q, d) + question_mask (B, Q)   from the question encoder
    choice_emb      (B, N, d) + choice_mask (B, N)     from the ChoiceEncoder (N is dynamic)

and returns one logit per choice (+ value, scalar score and the CLS embedding).
It knows nothing about Snake, text, bytes or tokenizers.

    [CLS] + question tokens + state tokens ──► Transformer ──► context H
                                                                  │
    scorer "cross": choices cross-attend H, then self-attend each other ─► Linear ─► logit_i
    scorer "bi":    logit_i = MLP([h, c_i, h*c_i])  (light cross scorer on pooled h = CLS)
    scorer "fixed_head": Linear(CLS) -> max_choices logits (ablation: ignores choice meaning)

"cross" and "bi" are permutation-equivariant over choices (no positional information).
"""
from __future__ import annotations

import torch
from torch import nn

from .layers import DecisionBlock, EncoderBlock

NEG = -1e9
SCORERS = ("cross", "bi", "fixed_head")


class MiniJEVCore(nn.Module):
    def __init__(self, d: int, n_layers: int, n_heads: int = 4, mlp_ratio: float = 4.0, dropout: float = 0.0,
                 n_decision_layers: int = 1, scorer: str = "cross", max_choices: int = 16):
        super().__init__()
        if scorer not in SCORERS:
            raise ValueError(f"scorer must be one of {SCORERS}, got {scorer!r}")
        self.d, self.scorer = d, scorer
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.type_emb = nn.Parameter(torch.zeros(3, d))  # cls / question / state
        nn.init.normal_(self.type_emb, std=0.02)
        self.blocks = nn.ModuleList(EncoderBlock(d, n_heads, mlp_ratio, dropout) for _ in range(n_layers))
        self.ln_ctx = nn.LayerNorm(d)
        if scorer == "cross":
            self.decision = nn.ModuleList(DecisionBlock(d, n_heads, mlp_ratio, dropout)
                                          for _ in range(n_decision_layers))
            self.choice_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 1))
        elif scorer == "bi":
            self.ln_h, self.ln_c = nn.LayerNorm(d), nn.LayerNorm(d)
            self.choice_head = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Linear(d, 1))
        else:
            self.choice_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, max_choices))
        self.value_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
        # continuous 0..1 score for `model.score(state, question)` (trained with BCE when rows carry `score`)
        self.scalar_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 1))

    def encode_context(self, state_tokens, state_mask, question_tokens, question_mask):
        B = question_tokens.shape[0]
        parts = [self.cls.expand(B, -1, -1) + self.type_emb[0], question_tokens + self.type_emb[1]]
        masks = [torch.ones(B, 1, dtype=torch.bool, device=question_tokens.device), question_mask]
        if state_tokens is not None:
            parts.append(state_tokens + self.type_emb[2])
            masks.append(state_mask)
        x, mask = torch.cat(parts, 1), torch.cat(masks, 1)
        for blk in self.blocks:
            x = blk(x, key_mask=mask)
        return self.ln_ctx(x), mask

    def score_choices(self, ctx, ctx_mask, choice_emb, choice_mask):
        if self.scorer == "cross":
            c = choice_emb
            for blk in self.decision:
                c = blk(c, ctx, ctx_mask, choice_mask)
            logits = self.choice_head(c).squeeze(-1)
        elif self.scorer == "bi":
            h = self.ln_h(ctx[:, :1]).expand(-1, choice_emb.shape[1], -1)
            c = self.ln_c(choice_emb)
            logits = self.choice_head(torch.cat([h, c, h * c], -1)).squeeze(-1)
        else:
            logits = self.choice_head(ctx[:, 0])[:, : choice_emb.shape[1]]
        return logits.float().masked_fill(~choice_mask, NEG)

    def forward(self, state_tokens, state_mask, question_tokens, question_mask, choice_emb, choice_mask):
        ctx, ctx_mask = self.encode_context(state_tokens, state_mask, question_tokens, question_mask)
        cls = ctx[:, 0]
        return {"logits": self.score_choices(ctx, ctx_mask, choice_emb, choice_mask),
                "value": self.value_head(cls).squeeze(-1).float(),
                "score": self.scalar_head(cls).squeeze(-1).float(),  # logit; sigmoid -> 0..1
                "embedding": cls, "choice_mask": choice_mask}