"""MIMIR-1 as a transformers model: load with `AutoModel`, forward on encoded batches.

Module and parameter names are the training checkpoint's state-dict keys. Renaming
one breaks `from_pretrained` for everybody who downloaded the release, so the head
below keeps the training code's attribute names exactly (`projection`, `roles`,
`typed_values`, `state_memory`, `candidate_processor`, `reasoner` with `.cell`,
`readout` with `utility`/`abstain`/`ordinal_score`/`threshold`/`histogram`), and the
encoder is a stock `ModernBertModel` under `self.encoder` for zero remapping.

The computation is the same one the shipped `model.onnx` already carries, so this
file ports the head's forward only: no loss, no training machinery. Text handling,
calibration, the server and MCP stay in the `mimir-decisions` package; this path is load
plus forward on the 21 flat named tensors the graph reads, returning every depth's
readouts, evidence and workspace.

The typed-value quantile edges live in the checkpoint buffers, so the constructor
takes them as optional: without them the buffers are zeros of the right shape and
`load_state_dict` fills the real values. A forward before loading is meaningless.

The encoder forward replicates the training encoder call exactly (embeddings, per
layer-kind rotary, direct layer calls with prepared dense masks, typed-value
injection at `injection_layer`, final norm). SDPA only; tested transformers==5.17.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any

import numpy as np
import numpy.typing as npt
import torch
from torch import nn
from transformers import ModernBertModel, PreTrainedModel
from transformers.utils.generic import ModelOutput

try:
    from .configuration_mimir import DECISION_TYPES, MimirConfig
except ImportError:
    # Loaded as remote code, where the subfolder is on sys.path but is no package.
    from configuration_mimir import (  # type: ignore[import-not-found,no-redef]
        DECISION_TYPES,
        MimirConfig,
    )

STATE_ROLE = 0
QUESTION_ROLE = 1
CANDIDATE_ROLE = 2

FULL_ATTENTION = "full_attention"
SLIDING_ATTENTION = "sliding_attention"

NUMBER = 0
DATE = 1
CYCLES = 4
SECONDS_PER_DAY = 86_400
FEBRUARY = 2
EPOCH_WEEKDAY = 3

RELATIONS = ("supports", "contradicts", "irrelevant")

INPUTS = (
    "chunk_ids",
    "chunk_mask",
    "chunk_record",
    "question_length",
    "key_chunk",
    "key_position",
    "key_mask",
    "candidate_ids",
    "candidate_mask",
    "candidate_text_mask",
    "candidate_index",
    "candidate_present",
    "decision_type",
    "typed_kind",
    "typed_number",
    "typed_year",
    "typed_month",
    "typed_day",
    "typed_second",
    "typed_chunk",
    "typed_position",
)
"""The forward's named inputs, in order: the same 21 the ONNX graph reads."""


class Norm(nn.RMSNorm):
    """RMSNorm in float32, returned in the input's dtype."""

    def forward(self, states: torch.Tensor) -> torch.Tensor:
        normed: torch.Tensor = super().forward(states.float())
        return normed.to(states.dtype)


@dataclass(frozen=True, slots=True)
class Memory:
    """Keys (normalised) and values over [batch, heads, length, size], with the key mask."""

    key: torch.Tensor
    value: torch.Tensor
    mask: torch.Tensor

    def extend(self, other: Memory) -> Memory:
        return Memory(
            torch.cat([self.key, other.key], dim=2),
            torch.cat([self.value, other.value], dim=2),
            torch.cat([self.mask, other.mask], dim=1),
        )


def masked_attention(query: torch.Tensor, memory: Memory) -> torch.Tensor:
    """Fused attention; a row with no valid key reads zeros instead of NaN."""
    bias = torch.zeros(memory.mask.shape, dtype=query.dtype, device=query.device)
    bias = bias.masked_fill(~memory.mask, torch.finfo(query.dtype).min)[:, None, None, :]
    read: torch.Tensor = nn.functional.scaled_dot_product_attention(
        query, memory.key, memory.value.to(query.dtype), attn_mask=bias
    )
    return read * memory.mask.any(dim=1)[:, None, None, None]


def attention_weights(query: torch.Tensor, memory: Memory) -> torch.Tensor:
    """The same attention's weights in float32, zero on masked keys and on keyless rows."""
    logits = (query @ memory.key.transpose(-1, -2)).float() / math.sqrt(query.shape[-1])
    mask = memory.mask[:, None, None, :]
    logits = logits.masked_fill(~mask, torch.finfo(logits.dtype).min)
    return logits.softmax(dim=-1) * mask.any(dim=-1, keepdim=True)


class Attention(nn.Module):
    """Multi-head attention from queries onto a memory, with RMSNorm on queries and keys."""

    def __init__(self, config: MimirConfig, *, zero_output: bool) -> None:
        super().__init__()
        self.heads = config.heads
        self.head_size = config.head_size
        self.query = nn.Linear(config.width, config.width, bias=False)
        self.key_value = nn.Linear(config.width, 2 * config.width, bias=False)
        self.output = nn.Linear(config.width, config.width, bias=False)
        self.query_norm = Norm(config.head_size)
        self.key_norm = Norm(config.head_size)
        if zero_output:
            nn.init.zeros_(self.output.weight)

    def _split(self, states: torch.Tensor) -> torch.Tensor:
        batch, length, _ = states.shape
        return states.view(batch, length, self.heads, self.head_size).transpose(1, 2)

    def _merge(self, read: torch.Tensor) -> torch.Tensor:
        batch, _, length, _ = read.shape
        width = self.heads * self.head_size
        merged: torch.Tensor = self.output(read.transpose(1, 2).reshape(batch, length, width))
        return merged

    def memory(self, states: torch.Tensor, mask: torch.Tensor) -> Memory:
        key, value = self.key_value(states).chunk(2, dim=-1)
        return Memory(self.key_norm(self._split(key)), self._split(value), mask)

    def queries(self, states: torch.Tensor) -> torch.Tensor:
        split: torch.Tensor = self.query_norm(self._split(self.query(states)))
        return split

    def read(self, states: torch.Tensor, memory: Memory) -> torch.Tensor:
        return self._merge(masked_attention(self.queries(states), memory))

    def read_with_weights(
        self, states: torch.Tensor, memory: Memory
    ) -> tuple[torch.Tensor, torch.Tensor]:
        weights = attention_weights(self.queries(states), memory)
        return self._merge(weights.to(memory.value.dtype) @ memory.value), weights


class AttentionBlock(nn.Module):
    """states + Attention(norm(states), norm(memory)); self-attention when memory is states."""

    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        self.query_norm = Norm(config.width)
        self.memory_norm = Norm(config.width)
        self.attention = Attention(config, zero_output=True)

    def memory(self, states: torch.Tensor, mask: torch.Tensor) -> Memory:
        return self.attention.memory(self.memory_norm(states), mask)

    def read(self, states: torch.Tensor, memory: Memory) -> torch.Tensor:
        return states + self.attention.read(self.query_norm(states), memory)

    def forward(
        self, states: torch.Tensor, memory: torch.Tensor, memory_mask: torch.Tensor
    ) -> torch.Tensor:
        return self.read(states, self.memory(memory, memory_mask))


class FeedForwardBlock(nn.Module):
    """states + SwiGLU(norm(states))."""

    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        self.norm = Norm(config.width)
        self.gate_up = nn.Linear(config.width, 2 * config.feedforward_hidden, bias=False)
        self.down = nn.Linear(config.feedforward_hidden, config.width, bias=False)
        nn.init.zeros_(self.down.weight)

    def forward(self, states: torch.Tensor) -> torch.Tensor:
        gate, up = self.gate_up(self.norm(states)).chunk(2, dim=-1)
        update: torch.Tensor = self.down(nn.functional.silu(gate) * up)
        return states + update


@dataclass(frozen=True, slots=True)
class TypedValues:
    kind: torch.Tensor
    number: torch.Tensor
    year: torch.Tensor
    month: torch.Tensor
    day: torch.Tensor
    second: torch.Tensor


def signed_log(values: torch.Tensor) -> torch.Tensor:
    return values.sign() * values.abs().log1p()


def days_from_civil(year: torch.Tensor, month: torch.Tensor, day: torch.Tensor) -> torch.Tensor:
    """Days since 1970-01-01 in the proleptic Gregorian calendar (Hinnant's algorithm)."""
    shifted = year - (month <= FEBRUARY).long()
    era = torch.div(shifted, 400, rounding_mode="floor")
    year_of_era = shifted - era * 400
    month_index = (month + 9) % 12
    day_of_year = torch.div(153 * month_index + 2, 5, rounding_mode="floor") + day - 1
    day_of_era = (
        year_of_era * 365
        + torch.div(year_of_era, 4, rounding_mode="floor")
        - torch.div(year_of_era, 100, rounding_mode="floor")
        + day_of_year
    )
    return era * 146_097 + day_of_era - 719_468


def piecewise_linear(values: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
    """[n] values over T+1 edges -> [n, T]; the first bin is open below, the last above."""
    low, high = edges[:-1], edges[1:]
    width = high - low
    safe = torch.where(width > 0, width, torch.ones_like(width))
    fraction = (values[:, None] - low) / safe
    fraction = torch.where(width > 0, fraction, (values[:, None] >= high).to(values.dtype))
    bins = edges.shape[0] - 1
    first = torch.arange(bins, device=values.device) == 0
    last = torch.arange(bins, device=values.device) == bins - 1
    fraction = torch.where(first, fraction, fraction.clamp(min=0))
    return torch.where(last, fraction, fraction.clamp(max=1))


class TypedValueEncoder(nn.Module):
    """Number bins, day bins and 4 date cycles -> MLP -> encoder width; zero at initialisation."""

    def __init__(
        self,
        config: MimirConfig,
        number_edges: npt.NDArray[np.float64] | None = None,
        day_edges: npt.NDArray[np.float64] | None = None,
    ) -> None:
        super().__init__()
        for name, edges in (("number_edges", number_edges), ("day_edges", day_edges)):
            if edges is None:
                continue
            if edges.shape != (config.value_bins + 1,) or np.any(np.diff(edges) < 0):
                message = f"{name} of shape {edges.shape} for {config.value_bins} sorted bins"
                raise ValueError(message)
        missing = np.zeros(config.value_bins + 1, dtype=np.float64)
        self.number_edges: torch.Tensor
        self.day_edges: torch.Tensor
        self.register_buffer(
            "number_edges",
            torch.from_numpy(number_edges if number_edges is not None else missing).float(),
        )
        self.register_buffer(
            "day_edges",
            torch.from_numpy(day_edges if day_edges is not None else missing).float(),
        )
        features = 2 * config.value_bins + 2 * CYCLES
        self.hidden = nn.Linear(features, config.width, bias=False)
        self.output = nn.Linear(config.width, config.encoder_width, bias=False)
        nn.init.zeros_(self.output.weight)

    def features(self, values: TypedValues) -> torch.Tensor:
        is_number = values.kind == NUMBER
        is_date = values.kind == DATE
        number = torch.nan_to_num(values.number.float(), nan=0.0, posinf=0.0, neginf=0.0)
        number_part = piecewise_linear(signed_log(number), self.number_edges)
        days = days_from_civil(values.year, values.month.clamp(min=1), values.day.clamp(min=1))
        day_part = piecewise_linear(days.float(), self.day_edges)
        weekday = (days + EPOCH_WEEKDAY) % 7
        angles = (
            2
            * math.pi
            * torch.stack(
                [
                    (values.month.float() - 1) / 12,
                    (values.day.float() - 1) / 31,
                    weekday.float() / 7,
                    values.second.float() / SECONDS_PER_DAY,
                ],
                dim=-1,
            )
        )
        cycles = torch.cat([angles.sin(), angles.cos()], dim=-1)
        return torch.cat(
            [
                number_part * is_number[:, None],
                day_part * is_date[:, None],
                cycles * is_date[:, None],
            ],
            dim=-1,
        )

    def forward(self, values: TypedValues) -> torch.Tensor:
        encoded: torch.Tensor = self.output(nn.functional.silu(self.hidden(self.features(values))))
        return encoded


class StateMemory(nn.Module):
    """Compress a record's chunk tokens into question-conditioned latents Z."""

    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        self.latents = nn.Parameter(torch.empty(config.latents, config.width))
        nn.init.orthogonal_(self.latents)
        self.question = nn.Linear(config.width, config.width, bias=False)
        self.read = AttentionBlock(config)
        self.read_feedforward = FeedForwardBlock(config)
        self.mix = AttentionBlock(config)
        self.mix_feedforward = FeedForwardBlock(config)

    def forward(
        self, question: torch.Tensor, tokens: torch.Tensor, token_mask: torch.Tensor
    ) -> torch.Tensor:
        latents: torch.Tensor = self.latents + self.question(question)[:, None]
        latents = self.read_feedforward(self.read(latents, tokens, token_mask))
        every = torch.ones(latents.shape[:2], dtype=torch.bool, device=latents.device)
        memory: torch.Tensor = self.mix_feedforward(self.mix(latents, latents, every))
        return memory


class CandidateProcessor(nn.Module):
    """Candidates as a set: each reads the question, then all attend to all, with no positions."""

    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        self.read_question = AttentionBlock(config)
        self.interact = AttentionBlock(config)
        self.feedforward = FeedForwardBlock(config)

    def forward(
        self,
        candidates: torch.Tensor,
        candidate_mask: torch.Tensor,
        question: torch.Tensor,
        question_mask: torch.Tensor,
    ) -> torch.Tensor:
        candidates = self.read_question(candidates, question, question_mask)
        candidates = self.interact(candidates, candidates, candidate_mask)
        refined: torch.Tensor = self.feedforward(candidates)
        return refined


@dataclass(frozen=True, slots=True)
class Context:
    """The fixed memories every iteration reads, already projected."""

    recall: Memory
    tokens: Memory
    candidate_state: Memory
    candidate_mask: torch.Tensor


@dataclass(frozen=True, slots=True)
class Iteration:
    workspace: torch.Tensor
    candidates: torch.Tensor
    evidence: torch.Tensor
    gap_weights: torch.Tensor


def _every(states: torch.Tensor) -> torch.Tensor:
    return torch.ones(states.shape[:2], dtype=torch.bool, device=states.device)


class ReasonerCell(nn.Module):
    """The shared reasoning cell over workspace slots, applied once per iteration."""

    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        self.recall = AttentionBlock(config)
        self.slot_norm = Norm(config.width)
        self.token_norm = Norm(config.width)
        self.evidence_norm = Norm(config.width)
        self.gap = Attention(config, zero_output=False)
        self.relation = nn.Linear(3 * config.width, len(RELATIONS), bias=False)
        nn.init.zeros_(self.relation.weight)
        self.support = nn.Linear(config.width, config.width, bias=False)
        self.contradict = nn.Linear(config.width, config.width, bias=False)
        nn.init.zeros_(self.support.weight)
        nn.init.zeros_(self.contradict.weight)
        self.workspace_feedforward = FeedForwardBlock(config)
        self.read_candidates = AttentionBlock(config)
        self.candidates_read = AttentionBlock(config)
        self.interact = AttentionBlock(config)
        self.candidate_feedforward = FeedForwardBlock(config)

    def context(
        self,
        memory: torch.Tensor,
        question: torch.Tensor,
        question_mask: torch.Tensor,
        tokens: torch.Tensor,
        token_mask: torch.Tensor,
        candidate_mask: torch.Tensor,
    ) -> Context:
        recall = torch.cat([memory, question], dim=1)
        state = torch.cat([memory, tokens], dim=1)
        return Context(
            recall=self.recall.memory(recall, torch.cat([_every(memory), question_mask], dim=1)),
            tokens=self.gap.memory(self.token_norm(tokens), token_mask),
            candidate_state=self.candidates_read.memory(
                state, torch.cat([_every(memory), token_mask], dim=1)
            ),
            candidate_mask=candidate_mask,
        )

    def forward(
        self,
        workspace: torch.Tensor,
        candidates: torch.Tensor,
        evidence: torch.Tensor,
        context: Context,
    ) -> Iteration:
        own = torch.cat([workspace, evidence], dim=1)
        workspace = self.recall.read(
            workspace, self.recall.memory(own, _every(own)).extend(context.recall)
        )
        query = self.slot_norm(workspace)
        found, gap_weights = self.gap.read_with_weights(query, context.tokens)
        found_norm = self.evidence_norm(found)
        relation = self.relation(torch.cat([query, found_norm, query * found_norm], dim=-1))
        weights = relation.float().softmax(dim=-1).to(workspace.dtype)
        workspace = (
            workspace
            + weights[..., 0:1] * self.support(found_norm)
            + weights[..., 1:2] * self.contradict(found_norm)
        )
        evidence = torch.cat([evidence, found], dim=1)
        workspace = self.workspace_feedforward(workspace)
        workspace = self.read_candidates(workspace, candidates, context.candidate_mask)
        own = torch.cat([workspace, evidence], dim=1)
        candidates = self.candidates_read.read(
            candidates,
            self.candidates_read.memory(own, _every(own)).extend(context.candidate_state),
        )
        candidates = self.interact(candidates, candidates, context.candidate_mask)
        candidates = self.candidate_feedforward(candidates)
        return Iteration(workspace, candidates, evidence, gap_weights)


class Reasoner(nn.Module):
    """Workspace slots seeded by the schema and the question, then the cell once per iteration."""

    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        self.iterations = config.iterations
        self.slots = nn.Parameter(torch.empty(config.slots, config.width))
        nn.init.orthogonal_(self.slots)
        self.schemas = nn.Embedding(len(DECISION_TYPES), config.width)
        self.question = nn.Linear(config.width, config.width, bias=False)
        self.cell = ReasonerCell(config)

    def forward(
        self,
        question_mean: torch.Tensor,
        decision_type: torch.Tensor,
        candidates: torch.Tensor,
        context: Context,
    ) -> list[Iteration]:
        seed = self.schemas(decision_type) + self.question(question_mean)
        workspace = self.slots + seed[:, None]
        evidence = workspace.new_zeros((workspace.shape[0], 0, workspace.shape[2]))
        iterations: list[Iteration] = []
        for _ in range(self.iterations):
            step = self.cell(workspace, candidates, evidence, context)
            iterations.append(step)
            workspace, candidates, evidence = step.workspace, step.candidates, step.evidence
        return iterations


@dataclass(frozen=True, slots=True)
class Readout:
    utilities: torch.Tensor
    abstain: torch.Tensor
    ordinal_score: torch.Tensor
    thresholds: torch.Tensor
    histogram_logits: torch.Tensor


def readout_mlp(inputs: int, width: int, outputs: int) -> nn.Sequential:
    return nn.Sequential(
        Norm(inputs),
        nn.Linear(inputs, width, bias=False),
        nn.SiLU(),
        nn.Linear(width, outputs),
    )


class DecisionReadout(nn.Module):
    def __init__(self, config: MimirConfig) -> None:
        super().__init__()
        width = config.width
        self.utility = readout_mlp(width, width, 1)
        self.abstain = readout_mlp(width, width, 1)
        self.ordinal_score = readout_mlp(width, width, 1)
        self.threshold = readout_mlp(2 * width, width, 1)
        self.histogram = readout_mlp(width, width, config.histogram_bins)

    def forward(self, pooled: torch.Tensor, candidates: torch.Tensor) -> Readout:
        """`pooled` is the workspace averaged over its slots."""
        pairs = torch.cat([candidates[:, :-1], candidates[:, 1:]], dim=-1)
        raw = self.threshold(pairs).squeeze(-1)
        steps = nn.functional.softplus(raw[:, 1:])
        thresholds = torch.cat([raw[:, :1], raw[:, :1] + steps.cumsum(dim=1)], dim=1)
        return Readout(
            utilities=self.utility(candidates).squeeze(-1),
            abstain=self.abstain(pooled).squeeze(-1),
            ordinal_score=self.ordinal_score(pooled).squeeze(-1),
            thresholds=thresholds,
            histogram_logits=self.histogram(pooled),
        )


@dataclass(frozen=True, slots=True)
class HeadInputs:
    """Encoder states per record: state keys, question tokens and pooled candidates."""

    tokens: torch.Tensor
    token_mask: torch.Tensor
    question: torch.Tensor
    question_mask: torch.Tensor
    candidates: torch.Tensor
    candidate_mask: torch.Tensor
    decision_type: torch.Tensor


@dataclass(frozen=True, slots=True)
class DepthOutput:
    """One depth's decision; `workspace` is the slot-averaged workspace the readout reads."""

    readout: Readout
    evidence_scores: torch.Tensor
    workspace: torch.Tensor


class DecisionHead(nn.Module):
    def __init__(
        self,
        config: MimirConfig,
        number_edges: npt.NDArray[np.float64] | None = None,
        day_edges: npt.NDArray[np.float64] | None = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.projection = nn.Linear(config.encoder_width, config.width, bias=False)
        self.roles = nn.Embedding(3, config.width)
        self.typed_values = TypedValueEncoder(config, number_edges, day_edges)
        self.state_memory = StateMemory(config)
        self.candidate_processor = CandidateProcessor(config)
        self.reasoner = Reasoner(config)
        self.readout = DecisionReadout(config)

    def _project(self, states: torch.Tensor, role: int) -> torch.Tensor:
        projected: torch.Tensor = self.projection(states)
        return projected + self.roles.weight[role]

    def forward(self, inputs: HeadInputs) -> list[DepthOutput]:
        tokens = self._project(inputs.tokens, STATE_ROLE)
        question = self._project(inputs.question, QUESTION_ROLE)
        candidates = self._project(inputs.candidates, CANDIDATE_ROLE)
        weight = inputs.question_mask[..., None].to(question.dtype)
        question_mean = (question * weight).sum(dim=1) / weight.sum(dim=1).clamp(min=1)
        memory = self.state_memory(question_mean, tokens, inputs.token_mask)
        candidates = self.candidate_processor(
            candidates, inputs.candidate_mask, question, inputs.question_mask
        )
        context = self.reasoner.cell.context(
            memory,
            question,
            inputs.question_mask,
            tokens,
            inputs.token_mask,
            inputs.candidate_mask,
        )
        iterations = self.reasoner(question_mean, inputs.decision_type, candidates, context)
        outputs: list[DepthOutput] = []
        for step in iterations:
            pooled = step.workspace.mean(dim=1)
            outputs.append(
                DepthOutput(
                    readout=self.readout(pooled, step.candidates),
                    evidence_scores=step.gap_weights.mean(dim=(1, 2)),
                    workspace=pooled,
                )
            )
        return outputs


def padded_dense_masks(
    attention_mask: torch.Tensor, sliding_window: int
) -> dict[str, torch.Tensor]:
    positions = torch.arange(attention_mask.shape[1], device=attention_mask.device)
    keys = attention_mask.bool()[:, None, None, :]
    near = (positions[:, None] - positions[None, :]).abs() <= sliding_window
    return {FULL_ATTENTION: keys, SLIDING_ATTENTION: keys & near}


def _pool_candidates(hidden: torch.Tensor, text_mask: torch.Tensor) -> torch.Tensor:
    """Each candidate row's mean state over its own text tokens."""
    weight = text_mask[..., None].to(hidden.dtype)
    return (hidden * weight).sum(dim=1) / weight.sum(dim=1).clamp(min=1)


def _mean_question(
    question: torch.Tensor, chunk_record: torch.Tensor, question_length: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Question positions 1.. of every chunk, averaged over the chunks of each record."""
    positions = torch.arange(question.shape[1], device=question.device)
    per_chunk = (positions < question_length[:, None]).to(question.dtype)
    records = torch.arange(chunk_record.max().item() + 1, device=question.device)
    owner = (chunk_record[None, :] == records[:, None]).to(question.dtype)
    summed = torch.einsum("rc,cqd->rqd", owner, question * per_chunk[..., None])
    counts = owner @ per_chunk
    return summed / counts.clamp(min=1)[..., None], counts > 0


@dataclass
class MimirModelOutput(ModelOutput):
    """Every depth's readouts stacked first, plus the final pooled workspace."""

    utilities: torch.Tensor | None = None
    thresholds: torch.Tensor | None = None
    abstain: torch.Tensor | None = None
    ordinal_score: torch.Tensor | None = None
    histogram_logits: torch.Tensor | None = None
    evidence: torch.Tensor | None = None
    workspace: torch.Tensor | None = None


class MimirModel(PreTrainedModel):
    """The decision model behind `AutoModel`: the stock encoder plus the head.

    The encoder lives under `self.encoder`, the head under `self.head`. Forward takes
    the 21 flat named tensors the release graph reads and returns every depth's
    readouts. The batch must hold at least one candidate row, and the question states
    are `question_cap` wide whatever the chunk width, exactly like the exported path
    the graph was cut from.
    """

    config_class = MimirConfig
    config: MimirConfig
    base_model_prefix = "mimir"
    _supports_sdpa = True

    def __init__(self, config: MimirConfig) -> None:
        super().__init__(config)
        self.encoder = ModernBertModel(config)
        self.head = DecisionHead(config)
        self.post_init()
        self.encoder.set_attn_implementation("sdpa")

    def _init_weights(self, module: nn.Module) -> None:
        # Constructors already initialise every head weight (default, zero or
        # orthogonal, exactly like training); the checkpoint overwrites them all.
        # The encoder ran its own init at construction; leave it alone here.
        if isinstance(module, Norm) and isinstance(module.weight, torch.nn.Parameter):
            nn.init.ones_(module.weight)

    def get_input_embeddings(self) -> nn.Module:
        return self.encoder.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.encoder.set_input_embeddings(value)

    def _encode(
        self,
        input_ids: torch.Tensor,
        mask: torch.Tensor,
        typed: TypedValues,
        typed_chunk: torch.Tensor,
        typed_position: torch.Tensor,
        *,
        inject: bool,
    ) -> torch.Tensor:
        positions = torch.arange(input_ids.shape[1], device=input_ids.device)
        positions = positions.expand(input_ids.shape[0], -1)
        masks = padded_dense_masks(mask.long(), int(self.config.sliding_window))
        injection: torch.Tensor | None = None
        injection_layer = -1
        if inject and typed.kind.numel():
            values = self.head.typed_values(typed)
            hidden = values.new_zeros((*input_ids.shape, values.shape[-1]))
            hidden = hidden.index_put((typed_chunk, typed_position), values, accumulate=True)
            injection = hidden
            injection_layer = self.config.injection_layer
        layer_types = self.encoder.config.layer_types
        if layer_types is None:
            message = "encoder config names no layer types"
            raise ValueError(message)
        hidden = self.encoder.embeddings(input_ids=input_ids)
        rotary = {
            kind: self.encoder.rotary_emb(hidden, positions, kind) for kind in set(layer_types)
        }
        for index, (layer, kind) in enumerate(zip(self.encoder.layers, layer_types, strict=True)):
            if injection is not None and index == injection_layer:
                hidden = hidden + injection.to(hidden.dtype)
            hidden = layer(hidden, attention_mask=masks[kind], position_embeddings=rotary[kind])
        output: torch.Tensor = self.encoder.final_norm(hidden)
        return output

    def forward(
        self,
        chunk_ids: torch.Tensor,
        chunk_mask: torch.Tensor,
        chunk_record: torch.Tensor,
        question_length: torch.Tensor,
        key_chunk: torch.Tensor,
        key_position: torch.Tensor,
        key_mask: torch.Tensor,
        candidate_ids: torch.Tensor,
        candidate_mask: torch.Tensor,
        candidate_text_mask: torch.Tensor,
        candidate_index: torch.Tensor,
        candidate_present: torch.Tensor,
        decision_type: torch.Tensor,
        typed_kind: torch.Tensor,
        typed_number: torch.Tensor,
        typed_year: torch.Tensor,
        typed_month: torch.Tensor,
        typed_day: torch.Tensor,
        typed_second: torch.Tensor,
        typed_chunk: torch.Tensor,
        typed_position: torch.Tensor,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> MimirModelOutput | tuple[torch.Tensor, ...]:
        if candidate_ids.shape[0] == 0:
            message = "the batch holds no candidate row; the exported path needs at least one"
            raise ValueError(message)
        question_cap = self.config.question_cap
        typed = TypedValues(
            kind=typed_kind,
            number=typed_number,
            year=typed_year,
            month=typed_month,
            day=typed_day,
            second=typed_second,
        )
        hidden = self._encode(
            chunk_ids, chunk_mask.bool(), typed, typed_chunk, typed_position, inject=True
        )
        question = nn.functional.pad(hidden, (0, 0, 0, question_cap + 1))[:, 1 : 1 + question_cap]
        candidates = self._encode(
            candidate_ids,
            candidate_mask.bool(),
            typed,
            typed_chunk,
            typed_position,
            inject=False,
        )
        pooled = _pool_candidates(candidates, candidate_text_mask.bool())
        question, question_mask = _mean_question(question, chunk_record, question_length)
        tokens = hidden[key_chunk, key_position]
        if pooled.shape[0]:
            selected = pooled[candidate_index]
        else:
            selected = pooled.new_zeros((*candidate_index.shape, pooled.shape[-1]))
        inputs = HeadInputs(
            tokens=tokens,
            token_mask=key_mask.bool(),
            question=question,
            question_mask=question_mask,
            candidates=selected,
            candidate_mask=candidate_present.bool(),
            decision_type=decision_type,
        )
        outputs = self.head(inputs)
        per_depth: dict[str, list[torch.Tensor]] = {
            "utilities": [],
            "thresholds": [],
            "abstain": [],
            "ordinal_score": [],
            "histogram_logits": [],
            "evidence": [],
        }
        for output in outputs:
            readout = output.readout
            per_depth["utilities"].append(readout.utilities.float())
            per_depth["thresholds"].append(readout.thresholds.float())
            per_depth["abstain"].append(readout.abstain.float())
            per_depth["ordinal_score"].append(readout.ordinal_score.float())
            per_depth["histogram_logits"].append(readout.histogram_logits.float())
            per_depth["evidence"].append(output.evidence_scores.float())
        stacked = {name: torch.stack(values) for name, values in per_depth.items()}
        utilities = stacked["utilities"]
        thresholds = stacked["thresholds"]
        abstain = stacked["abstain"]
        ordinal_score = stacked["ordinal_score"]
        histogram_logits = stacked["histogram_logits"]
        evidence = stacked["evidence"]
        workspace = outputs[-1].workspace.float()
        result = MimirModelOutput(
            utilities=utilities,
            thresholds=thresholds,
            abstain=abstain,
            ordinal_score=ordinal_score,
            histogram_logits=histogram_logits,
            evidence=evidence,
            workspace=workspace,
        )
        if return_dict is False:
            return (
                utilities,
                thresholds,
                abstain,
                ordinal_score,
                histogram_logits,
                evidence,
                workspace,
            )
        return result


__all__ = ["INPUTS", "MimirModel", "MimirModelOutput"]