"""Configuration dataclass for PRISME models (backbone + workspace + controller)."""
from __future__ import annotations

import dataclasses
import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml


@dataclass
class PrismeConfig:
    # --- backbone -----------------------------------------------------------
    name: str = "prisme_tiny"
    vocab_size: int = 32768
    d_model: int = 768
    n_layers: int = 12
    n_heads: int = 12
    n_kv_heads: int = 3
    ffn_width: int = 2048
    max_seq_len: int = 4096
    rope_theta: float = 10000.0
    norm_eps: float = 1e-5
    tie_embeddings: bool = True
    init_std: float = 0.02
    # --- workspace ----------------------------------------------------------
    use_workspace: bool = True
    w: int = 256                 # workspace width
    ffn_w: int = 768             # workspace FFN width
    n_slots: int = 16            # M
    n_editor_blocks: int = 2     # Lw (shared weights, applied sequentially)
    n_editor_heads: int = 4
    r_max: int = 8               # hard cap on actions per token
    eta_max: float = 0.25
    clip_c: float = 1.0
    clip_eps: float = 1e-6
    readout_zero_init: bool = True
    readout_init_std: float = 0.02  # used only when readout_zero_init is False (bootstrap ablation)
    readout_scale: float = 1.0   # fixed (non-learned) multiplier on U(...); set to the backbone residual rms at bootstrap (D9)
    # --- controller ---------------------------------------------------------
    controller_hidden: int = 0   # 0 -> ffn_w
    ctrl_input: str = "z"        # "z" (w-dim summary) or "h" (normalised backbone residual, d-dim; costs 2dw per token)
    ctrl_use_entropy: bool = False  # add the backbone's own predictive entropy at budget 0 (costs one head evaluation per token)
    lambda_cost: float = 0.0
    kappa_risk: float = 0.0
    # --- dense-workspace ablation (writes all slots) --------------------------
    dense_writes: bool = False

    @property
    def head_dim(self) -> int:
        assert self.d_model % self.n_heads == 0
        return self.d_model // self.n_heads

    @property
    def kv_dim(self) -> int:
        return self.n_kv_heads * self.head_dim

    @property
    def editor_head_dim(self) -> int:
        assert self.w % self.n_editor_heads == 0
        return self.w // self.n_editor_heads

    @property
    def n_pairs(self) -> int:
        return self.n_slots // 2

    @property
    def n_actions(self) -> int:
        """Nonterminal actions (singles + fixed pairs); STOP is index n_actions."""
        return self.n_slots + self.n_pairs

    @property
    def stop_action(self) -> int:
        return self.n_actions

    def validate(self) -> None:
        assert self.n_heads % self.n_kv_heads == 0, "GQA requires n_heads % n_kv_heads == 0"
        assert self.n_slots % 2 == 0, "fixed pairs require an even slot count"
        assert self.w % self.n_editor_heads == 0
        assert self.r_max >= 1

    def to_dict(self) -> dict[str, Any]:
        return dataclasses.asdict(self)

    def sha256(self) -> str:
        return hashlib.sha256(json.dumps(self.to_dict(), sort_keys=True).encode()).hexdigest()

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "PrismeConfig":
        known = {f.name for f in dataclasses.fields(cls)}
        unknown = set(d) - known
        if unknown:
            raise ValueError(f"unknown config keys: {sorted(unknown)}")
        cfg = cls(**d)
        cfg.validate()
        return cfg


def load_config(path: str | Path) -> PrismeConfig:
    with open(path) as f:
        raw = yaml.safe_load(f)
    model = raw.get("model", raw)
    return PrismeConfig.from_dict(model)