"""Hev decision head: projection, Hydra stack, pointer readout over option spans, backbone prior, temperatures."""
from __future__ import annotations

import math
from dataclasses import asdict, dataclass

import torch
from torch import nn

from .cache import Batch
from .hydra_torch import HydraBlock, RMSNormGated


@dataclass
class HevHeadConfig:
    d_backbone: int = 2816
    d_model: int = 512
    n_blocks: int = 3
    d_state: int = 64
    expand: int = 2
    headdim: int = 64
    d_conv: int = 7
    mlp_mult: int = 2
    chunk_size: int = 64
    n_types: int = 3

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "HevHeadConfig":
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


class HevHead(nn.Module):
    def __init__(self, cfg: HevHeadConfig):
        super().__init__()
        self.cfg = cfg
        self.in_proj = nn.Linear(cfg.d_backbone, cfg.d_model)
        self.seg_emb = nn.Embedding(2, cfg.d_model)
        self.in_norm = RMSNormGated(cfg.d_model)
        self.blocks = nn.ModuleList([
            HydraBlock(cfg.d_model, mlp_mult=cfg.mlp_mult, d_state=cfg.d_state, d_conv=cfg.d_conv,
                       expand=cfg.expand, headdim=cfg.headdim, chunk_size=cfg.chunk_size)
            for _ in range(cfg.n_blocks)
        ])
        self.final_norm = RMSNormGated(cfg.d_model)
        self.w_q = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.w_o = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        nn.init.normal_(self.w_q.weight, std=0.02)
        nn.init.zeros_(self.w_o.weight)  # pointer term is exactly 0 at init, so a fresh head starts at the backbone prior
        self.beta = nn.Parameter(torch.ones(cfg.n_types))
        self.log_temp = nn.Parameter(torch.zeros(cfg.n_types))

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def encode(self, feats: torch.Tensor, seg: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        h = self.in_proj(feats.float()) + self.seg_emb(seg)
        h = self.in_norm(h) * mask[..., None].to(h.dtype)
        for blk in self.blocks:
            h = blk(h, mask)
        return self.final_norm(h)

    def forward_flat(self, batch: Batch, use_temperature: bool = False) -> torch.Tensor:
        """Logits for every option of every decision in the batch as one flat tensor, in the batch's option order.

        This is the whole readout in a fixed number of tensor ops: a gather of the span rows, one index_add for
        the span means, one gather for the slot queries, and elementwise arithmetic on the flat option axis.
        The per-decision loops it replaced issued ~360 ops per decision and left a GPU idle (session 3).
        """
        h = self.encode(batch.feats, batch.seg, batch.mask)
        b, l, d = h.shape
        if bool((batch.span_len <= 0).any()):  # one host sync per batch; collate already rejects empty spans
            bad = int((batch.span_len <= 0).nonzero()[0])
            raise ValueError(f"empty span for decision {batch.decisions[int(batch.opt_dec[bad])]['qid']!r}")
        hf = h.reshape(b * l, d).float()
        sums = hf.new_zeros(batch.span_len.shape[0], d).index_add(0, batch.span_opt, hf.index_select(0, batch.span_rows))
        reps = sums / batch.span_len[:, None]
        q = self.w_q(h[batch.dec_b, batch.dec_slot].float())  # (n_dec, d_model)
        scale = 1.0 / math.sqrt(self.cfg.d_model)
        scores = (self.w_o(reps) * q[batch.opt_dec]).sum(-1).float() * scale
        logits = scores + self.beta[batch.dec_type][batch.opt_dec] * batch.opt_prior
        if use_temperature:
            logits = logits / torch.exp(self.log_temp)[batch.dec_type][batch.opt_dec]
        return logits

    def forward(self, batch: Batch, use_temperature: bool = False) -> list[torch.Tensor]:
        """One logit vector per decision: views into forward_flat's output."""
        return list(torch.split(self.forward_flat(batch, use_temperature), batch.counts))


class HevEnsemble(nn.Module):
    """One or more HevHeads with the same config, averaged in probability space.

    forward() returns, per decision, the log of the mean of the members' softmaxes, so every consumer that
    applies softmax to the head's output (evaluate.predict, HevForDecisions.decide) gets exactly the averaged
    distribution. Members keep their own fitted temperatures; the ensemble is not re-calibrated."""

    def __init__(self, members: list[HevHead]):
        super().__init__()
        if not members:
            raise ValueError("an ensemble needs at least one member")
        cfg = members[0].cfg
        if any(m.cfg != cfg for m in members):
            raise ValueError("every ensemble member must share one HevHeadConfig")
        self.members = nn.ModuleList(members)
        self.cfg = cfg

    def num_parameters(self) -> int:
        return sum(m.num_parameters() for m in self.members)

    def forward(self, batch: Batch, use_temperature: bool = False) -> list[torch.Tensor]:
        outs = [m(batch, use_temperature=use_temperature) for m in self.members]
        result = []
        for i in range(len(outs[0])):
            probs = torch.stack([torch.softmax(o[i].float(), 0) for o in outs]).mean(0)
            result.append(torch.log(probs.clamp_min(1e-12)))
        return result