"""Typed configuration objects and deterministic hashing.

Every experiment is described by a pydantic model; ``config_hash`` is the SHA-256 of the
canonical JSON dump so that identical configurations always produce identical run IDs.
"""
from __future__ import annotations

import hashlib
import json
from typing import Any, Literal

from pydantic import BaseModel, Field


def canonical_json(obj: Any) -> str:
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=True, default=str)


def config_hash(model: BaseModel, length: int = 12) -> str:
    return hashlib.sha256(canonical_json(model.model_dump()).encode()).hexdigest()[:length]


class OpSeqTaskConfig(BaseModel):
    """Micro-world A: compositional operator sequences over Z_p (see datasets/opseq.py)."""

    family: Literal["opseq"] = "opseq"
    modulus: int = 11
    operators: list[str] = Field(default_factory=lambda: ["zib", "zob", "qua", "lar", "neg", "flip", "dro", "mik"])
    min_len: int = 1
    max_len: int = 5
    held_out_pairs: list[tuple[str, str]] = Field(default_factory=lambda: [("zob", "lar"), ("qua", "neg")])
    """Adjacent operator pairs never seen in training (compositional held-out)."""
    excluded_operators: list[str] = Field(default_factory=list)
    """Operators removed from a model's training distribution (used to create recipient weakness)."""
    trace_stride: int = 1
    """Emit every ``trace_stride``-th intermediate result (plus the final one). 1 = full scratchpad;
    larger values force multi-step composition inside a single prediction (difficulty knob)."""


class FSATaskConfig(BaseModel):
    """Micro-world B: run a fixed random DFA over an invented alphabet and output the final state."""

    family: Literal["fsa"] = "fsa"
    n_states: int = 4
    alphabet: list[str] = Field(default_factory=lambda: ["ka", "mo", "ti"])
    min_len: int = 1
    max_len: int = 8
    automaton_seed: int = 7


class MorphTaskConfig(BaseModel):
    """Micro-world C: context-dependent morphological rule on pseudo-words."""

    family: Literal["morph"] = "morph"
    min_stem: int = 3
    max_stem: int = 6
    rules: dict[str, str] = Field(default_factory=lambda: {"vowel": "-na", "consonant": "-ek", "nasal": "-ur"})


TaskConfig = OpSeqTaskConfig | FSATaskConfig | MorphTaskConfig


class ToyModelConfig(BaseModel):
    d_model: int = 128
    n_layers: int = 4
    n_heads: int = 4
    d_ff: int = 512
    max_seq_len: int = 32
    dropout: float = 0.0
    tie_embeddings: bool = False


class TrainConfig(BaseModel):
    seed: int = 0
    n_train: int = 40_000
    n_dev: int = 2_000
    batch_size: int = 256
    steps: int = 3_000
    lr: float = 1e-3
    weight_decay: float = 0.01
    warmup_steps: int = 100
    device: Literal["auto", "cpu", "mps", "cuda"] = "auto"
    eval_every: int = 500


class ModelRoleConfig(BaseModel):
    """A trained toy model with a role in the donor/recipient design."""

    role: Literal["donor", "recipient", "unrelated"]
    name: str
    model: ToyModelConfig = Field(default_factory=ToyModelConfig)
    train: TrainConfig = Field(default_factory=TrainConfig)
    task: OpSeqTaskConfig | FSATaskConfig | MorphTaskConfig = Field(default_factory=OpSeqTaskConfig)


class PilotConfig(BaseModel):
    experiment_id: str = "pilot_tier0_opseq"
    description: str = "Tier-0 pilot: donor/recipient gap, hook no-op test, within-model causal localization"
    seeds: list[int] = Field(default_factory=lambda: [0, 1, 2])
    donor: ModelRoleConfig
    recipient: ModelRoleConfig
    target_operator: str = "qua"
    n_eval: int = 1_000
    n_bootstrap: int = 2_000