"""Configuration for the OxMini model family."""

from __future__ import annotations

from dataclasses import asdict, dataclass, fields
import json
from pathlib import Path
from typing import Any

import yaml


@dataclass
class OxMiniConfig:
    """Serializable OxMini architecture configuration.

    ``use_mhc=False`` selects the hybrid-attention-only ablation. The separate
    :class:`BaselineGPTForCausalLM` uses the same width/depth fields.
    """

    vocab_size: int
    block_size: int = 256
    n_layer: int = 8
    n_embd: int = 128
    n_head: int = 4
    kda_to_mla_ratio: int = 3
    mla_latent_dim: int = 32
    hc_streams: int = 4
    use_sinkhorn_mhc: bool = False
    dropout: float = 0.1
    tie_embeddings: bool = True
    use_mhc: bool = True
    ffn_multiplier: int = 4
    rms_norm_eps: float = 1e-5
    bias: bool = False
    seed: int = 1337

    def __post_init__(self) -> None:
        if self.vocab_size <= 0:
            raise ValueError("vocab_size must be positive")
        if self.block_size <= 0 or self.n_layer <= 0 or self.n_embd <= 0:
            raise ValueError("block_size, n_layer, and n_embd must be positive")
        if self.n_head <= 0 or self.n_embd % self.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")
        if (self.n_embd // self.n_head) % 2:
            raise ValueError("head dimension must be even for RoPE")
        if not 0 < self.mla_latent_dim < self.n_embd:
            raise ValueError("mla_latent_dim must be between 0 and n_embd")
        if self.hc_streams <= 0:
            raise ValueError("hc_streams must be positive")
        if self.kda_to_mla_ratio <= 0:
            raise ValueError("kda_to_mla_ratio must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")
        if self.ffn_multiplier <= 0:
            raise ValueError("ffn_multiplier must be positive")

    @property
    def head_dim(self) -> int:
        return self.n_embd // self.n_head

    @property
    def layer_types(self) -> list[str]:
        """Return the repeated KDA:MLA pattern, with a final MLA on remainders."""

        # For the default ratio, one group is [KDA, KDA, KDA, MLA]. Placing a
        # global-attention layer at the end of an incomplete group avoids ending
        # a non-standard-depth model with only recurrent local-state updates.
        group = self.kda_to_mla_ratio + 1
        result: list[str] = []
        for index in range(self.n_layer):
            at_group_end = (index + 1) % group == 0
            is_final_remainder = index == self.n_layer - 1 and self.n_layer % group != 0
            result.append("mla" if at_group_end or is_final_remainder else "kda")
        return result

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

    @classmethod
    def from_dict(cls, values: dict[str, Any]) -> "OxMiniConfig":
        # Hub config files include descriptive keys such as ``architectures``
        # and ``model_type``. Ignoring unknown keys keeps the plain dataclass
        # loader compatible with those metadata-rich JSON files.
        allowed = {field.name for field in fields(cls)}
        return cls(**{key: value for key, value in values.items() if key in allowed})

    @classmethod
    def from_file(cls, path: str | Path) -> "OxMiniConfig":
        path = Path(path)
        with path.open("r", encoding="utf-8") as handle:
            if path.suffix.lower() == ".json":
                values = json.load(handle)
            else:
                values = yaml.safe_load(handle)
        if not isinstance(values, dict):
            raise ValueError(f"configuration in {path} must be a mapping")
        return cls.from_dict(values)

    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as handle:
            if path.suffix.lower() == ".json":
                json.dump(self.to_dict(), handle, indent=2, sort_keys=True)
                handle.write("\n")
            else:
                yaml.safe_dump(self.to_dict(), handle, sort_keys=False)