"""Run configuration: model shape and training hyperparameters, loaded from YAML."""

from __future__ import annotations

import dataclasses
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

import yaml


@dataclass
class ModelConfig:
    d_model: int
    n_layers: int
    n_heads: int
    ffn_hidden: int
    context_len: int
    vocab_size: int = 16384
    rope_theta: float = 10000.0
    norm_eps: float = 1e-5
    init_std: float = 0.02


@dataclass
class TrainConfig:
    peak_lr: float
    total_steps: int
    micro_batch_size: int
    grad_accum: int
    precision: Literal["fp32", "fp16"]
    data_dir: str
    mix: dict[str, float]
    max_epochs: dict[str, float] = field(default_factory=dict)
    warmup_steps: int = 500
    min_lr_ratio: float = 0.1
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    checkpoint_minutes: float = 20.0
    eval_every: int = 200
    eval_batches: int = 20
    seed: int = 1337


@dataclass
class RunConfig:
    name: str
    model: ModelConfig
    train: TrainConfig


def _build(cls, data: dict, section: str):
    known = {f.name for f in dataclasses.fields(cls)}
    unknown = set(data) - known
    if unknown:
        raise ValueError(f"unknown keys in '{section}': {sorted(unknown)}")
    return cls(**data)


def load_run_config(path) -> RunConfig:
    data = yaml.safe_load(Path(path).read_text())
    unknown = set(data) - {"name", "model", "train"}
    if unknown:
        raise ValueError(f"unknown top-level keys: {sorted(unknown)}")
    return RunConfig(
        name=data["name"],
        model=_build(ModelConfig, data["model"], "model"),
        train=_build(TrainConfig, data["train"], "train"),
    )