from __future__ import annotations
from dataclasses import asdict, dataclass, fields
import json
import math
from pathlib import Path


@dataclass(frozen=True)
class ModelConfig:
    vocab_size: int = 32000
    layers: int = 12
    hidden_size: int = 768
    heads: int = 12
    intermediate_size: int = 2048
    context_length: int = 2048
    rope_theta: float = 10000.0
    rms_epsilon: float = 1e-5
    dropout: float = 0.0
    gradient_checkpointing: bool = False

    def __post_init__(self):
        dims=(self.vocab_size, self.layers, self.hidden_size, self.heads, self.intermediate_size, self.context_length)
        if any(type(value) is not int for value in dims) or min(dims) <= 0:
            raise ValueError("Model dimensions must be positive")
        if self.hidden_size % self.heads or (self.hidden_size // self.heads) % 2:
            raise ValueError("Head dimension must be an even integer")
        if not all(math.isfinite(value) for value in (self.dropout,self.rms_epsilon,self.rope_theta)) or not 0 <= self.dropout < 1 or self.rms_epsilon <= 0 or self.rope_theta <= 0:
            raise ValueError("Invalid dropout/RMS epsilon/RoPE theta")

    @property
    def parameter_count(self) -> int:
        h, f = self.hidden_size, self.intermediate_size
        return self.vocab_size*h + self.layers*(4*h*h + 3*h*f + 2*h) + h

    def to_dict(self):
        return asdict(self)


@dataclass(frozen=True)
class TrainConfig:
    batch_size: int = 2
    gradient_accumulation: int = 16
    max_steps: int = 1000
    schedule_total_steps: int | None = None
    learning_rate: float = 3e-4
    min_learning_rate_ratio: float = 0.1
    warmup_steps: int = 100
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    seed: int = 1337
    precision: str = "auto"
    checkpoint_every: int = 100
    eval_every: int = 100
    eval_batches: int = 10
    deterministic: bool = False
    max_full_checkpoints_per_session: int = 3
    checkpoint_interval_minutes: float = 0.0
    checkpoint_token_milestones: tuple[int, ...] = ()

    def __post_init__(self):
        for name in ("batch_size", "gradient_accumulation", "max_steps", "checkpoint_every", "eval_every", "eval_batches"):
            if type(getattr(self,name)) is not int or getattr(self, name) <= 0:
                raise ValueError(f"{name} must be positive")
        if self.schedule_total_steps is not None and (type(self.schedule_total_steps) is not int or self.schedule_total_steps < self.max_steps):
            raise ValueError("Explicit LR schedule horizon must cover this phase's max_steps")
        if not isinstance(self.checkpoint_token_milestones,(list,tuple)):
            raise ValueError("Checkpoint token milestones must be a list/tuple")
        milestones=tuple(self.checkpoint_token_milestones)
        if any(type(v) is not int or v<=0 for v in milestones) or tuple(sorted(set(milestones))) != milestones:
            raise ValueError("Checkpoint milestones must be strictly increasing positive integers")
        if len(milestones)>self.max_full_checkpoints_per_session-2:
            raise ValueError("Milestones must leave start and final checkpoint slots")
        object.__setattr__(self,"checkpoint_token_milestones",milestones)
        if type(self.seed) is not int or self.seed < 0:
            raise ValueError("Seed must be a nonnegative integer")
        if type(self.warmup_steps) is not int or not 0 <= self.warmup_steps <= self.max_steps:
            raise ValueError("warmup_steps must be between zero and max_steps")
        if not all(math.isfinite(value) for value in (self.learning_rate,self.grad_clip,self.weight_decay,self.min_learning_rate_ratio,self.beta1,self.beta2)):
            raise ValueError("Optimizer values must be finite")
        if self.learning_rate <= 0 or self.grad_clip <= 0 or self.weight_decay < 0:
            raise ValueError("Invalid optimizer parameters")
        if not 0 <= self.min_learning_rate_ratio <= 1 or not 0 <= self.beta1 < 1 or not 0 <= self.beta2 < 1:
            raise ValueError("Invalid optimizer ratios")
        if type(self.max_full_checkpoints_per_session) is not int or not 2 <= self.max_full_checkpoints_per_session <= 3:
            raise ValueError("Keep 2–3 full checkpoint slots per session")
        if not math.isfinite(self.checkpoint_interval_minutes) or self.checkpoint_interval_minutes < 0:
            raise ValueError("Checkpoint interval minutes must be finite and nonnegative (zero means adaptive)")
        if self.precision not in {"auto", "float32", "bfloat16", "float16"}:
            raise ValueError("Unsupported precision")

    def to_dict(self):
        return asdict(self)


def load_config(path: str | Path):
    with open(path) as f:
        obj = json.load(f)
    unknown = set(obj) - {"model", "training", "description", "provisional"}
    if unknown:
        raise ValueError(f"Unknown config keys: {sorted(unknown)}")
    return ModelConfig(**obj["model"]), TrainConfig(**obj["training"])