"""Pluto: a compact typed-decision model.

A bidirectional encoder reads the state, the typed question and every candidate
option. Each option is scored at its own marker token. A small permutation-
equivariant interaction stack compares options globally, then a scalar head
returns one logit per option. A second head predicts the teacher's confidence so
the runtime can abstain under a calibrated risk target.
"""

from __future__ import annotations

import json
import math
import os
from dataclasses import asdict, dataclass

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

from .schema import KIND_TO_ID

IGNORE_INDEX = -1


@dataclass
class PlutoArch:
    hidden: int = 768
    interaction_layers: int = 2
    interaction_heads: int = 8
    dropout: float = 0.05
    max_options: int = 192
    max_score_levels: int = 12


class PlutoForDecisions(nn.Module):
    def __init__(
        self,
        backbone_name: str | None = None,
        backbone_config: AutoConfig | None = None,
        arch: PlutoArch | None = None,
        vocab_size: int | None = None,
        freeze_layers: int = 0,
    ):
        super().__init__()
        self.arch = arch or PlutoArch()
        if backbone_config is None:
            self.backbone = AutoModel.from_pretrained(
                backbone_name, attn_implementation="sdpa"
            )
        else:
            self.backbone = AutoModel.from_config(
                backbone_config, attn_implementation="sdpa"
            )
        if vocab_size is not None and vocab_size != self.backbone.config.vocab_size:
            self.backbone.resize_token_embeddings(vocab_size)
        self.arch.hidden = self.backbone.config.hidden_size

        h = self.arch.hidden
        self.kind_embedding = nn.Embedding(len(KIND_TO_ID), h)
        self.score_position = nn.Embedding(self.arch.max_score_levels, h)
        self.option_in = nn.Sequential(nn.Linear(h, h), nn.LayerNorm(h))
        self.global_in = nn.Sequential(nn.Linear(h, h), nn.LayerNorm(h))
        layer = nn.TransformerEncoderLayer(
            d_model=h,
            nhead=self.arch.interaction_heads,
            dim_feedforward=2 * h,
            dropout=self.arch.dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.interaction = nn.TransformerEncoder(
            layer, num_layers=self.arch.interaction_layers, enable_nested_tensor=False
        )
        self.option_norm = nn.LayerNorm(h)
        self.scorer = nn.Sequential(
            nn.Linear(h, h),
            nn.GELU(),
            nn.Dropout(self.arch.dropout),
            nn.LayerNorm(h),
            nn.Linear(h, 1),
        )
        self.confidence_head = nn.Sequential(
            nn.Linear(h, h // 2), nn.GELU(), nn.Linear(h // 2, 1)
        )

        if freeze_layers > 0:
            for p in self.backbone.embeddings.parameters():
                p.requires_grad = False
            for layer_mod in self.backbone.layers[:freeze_layers]:
                for p in layer_mod.parameters():
                    p.requires_grad = False

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        marker_positions: torch.Tensor,
        n_options: torch.Tensor,
        kind_ids: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        hidden = self.backbone(
            input_ids=input_ids, attention_mask=attention_mask
        ).last_hidden_state

        b, nmax = marker_positions.shape
        h = hidden.shape[-1]
        safe_positions = marker_positions.clamp(min=0).unsqueeze(-1).expand(-1, -1, h)
        gathered = torch.gather(hidden, 1, safe_positions)

        option_ids = torch.arange(nmax, device=input_ids.device).unsqueeze(0)
        valid = option_ids < n_options.unsqueeze(1)

        kind_vec = self.kind_embedding(kind_ids)
        options = self.option_in(gathered) + kind_vec.unsqueeze(1)
        score_rows = kind_ids == KIND_TO_ID["score"]
        if score_rows.any():
            pos = torch.arange(nmax, device=input_ids.device).clamp(
                max=self.arch.max_score_levels - 1
            )
            pos_vec = self.score_position(pos).unsqueeze(0)
            options = options + pos_vec * score_rows.view(-1, 1, 1).to(options.dtype)

        global_vec = self.global_in(hidden[:, 0]) + kind_vec
        sequence = torch.cat([global_vec.unsqueeze(1), options], dim=1)
        pad_mask = torch.cat(
            [
                torch.zeros(b, 1, dtype=torch.bool, device=input_ids.device),
                ~valid,
            ],
            dim=1,
        )
        sequence = self.interaction(sequence, src_key_padding_mask=pad_mask)

        logits = self.scorer(sequence[:, 1:]).squeeze(-1).float()
        logits = logits.masked_fill(~valid, -1e9)
        correct_logit = self.confidence_head(sequence[:, 0]).squeeze(-1).float()
        return {
            "logits": logits,
            "correct_logit": correct_logit,
            "valid": valid,
        }


def decision_loss(
    outputs: dict[str, torch.Tensor],
    targets: torch.Tensor,
    kind_ids: torch.Tensor,
    weights: dict[str, float] | None = None,
) -> tuple[torch.Tensor, dict[str, float]]:
    weights = weights or {"kl": 1.0, "ce": 0.5, "rps": 0.5, "conf": 0.1}
    logits = outputs["logits"]
    valid = outputs["valid"]
    target = targets.clamp_min(0.0)
    target = target / target.sum(-1, keepdim=True).clamp_min(1e-9)
    log_probs = F.log_softmax(logits, dim=-1)
    row_valid = valid.any(-1)

    kind_w = torch.tensor([1.0, 0.8, 1.2], device=logits.device)[kind_ids]
    kl = -(target * log_probs).sum(-1)
    kl = (kl * kind_w)[row_valid].mean()

    argmax = target.argmax(-1)
    ce = F.nll_loss(log_probs, argmax, reduction="none")
    ce = (ce * kind_w)[row_valid].mean()

    score_rows = (kind_ids == KIND_TO_ID["score"]) & row_valid
    rps = torch.zeros((), device=logits.device)
    if score_rows.any():
        probs = log_probs.exp()
        cs = (probs.cumsum(-1) - target.cumsum(-1))[:, :-1]
        row_rps = cs.pow(2).mean(-1)
        rps = row_rps[score_rows].mean()

    correct_target = (logits.argmax(-1) == argmax).float()
    conf = F.binary_cross_entropy_with_logits(
        outputs["correct_logit"], correct_target, reduction="none"
    )
    conf = conf[row_valid].mean()

    total = (
        weights["kl"] * kl
        + weights["ce"] * ce
        + weights["rps"] * rps
        + weights["conf"] * conf
    )
    parts = {
        "kl": float(kl.detach()),
        "ce": float(ce.detach()),
        "rps": float(rps.detach()),
        "correct": float(conf.detach()),
    }
    return total, parts


def save_pluto(
    model: PlutoForDecisions,
    tokenizer,
    out_dir: str,
    extra: dict,
) -> None:
    os.makedirs(out_dir, exist_ok=True)
    state = {k: v.detach().cpu() for k, v in model.state_dict().items()}
    save_file(state, os.path.join(out_dir, "model.safetensors"))
    tokenizer.save_pretrained(out_dir)
    model.backbone.config.to_json_file(os.path.join(out_dir, "backbone_config.json"))
    config = {
        "format": "pluto-1",
        "arch": asdict(model.arch),
        "backbone_type": model.backbone.config.model_type,
        "vocab_size": model.backbone.config.vocab_size,
        **extra,
    }
    with open(os.path.join(out_dir, "pluto_config.json"), "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2)


def load_pluto(path: str, device: str | torch.device = "cpu") -> tuple[PlutoForDecisions, dict]:
    with open(os.path.join(path, "pluto_config.json"), "r", encoding="utf-8") as f:
        config = json.load(f)
    backbone_config = AutoConfig.for_model(
        **json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8"))
    )
    arch = PlutoArch(**config["arch"])
    model = PlutoForDecisions(
        backbone_config=backbone_config,
        arch=arch,
        vocab_size=config["vocab_size"],
    )
    state = load_file(os.path.join(path, "model.safetensors"))
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()
    return model, config