"""Experiment v1 configuration.

One experiment = one config file. Iteration 2 is a copy of this file with
different numbers, not a code change.

v1 is deliberately small (~4.3M parameters): the goal is a complete, legible
pipeline, not a strong model. It sits below TinyStories' coherence sweet spot
(~30M params), so some incoherence in the output is expected.
"""

from dataclasses import dataclass


@dataclass(frozen=True)
class ModelConfig:
    vocab_size: int = 4096       # matches the trained tokenizer
    context_length: int = 256    # tokens of history the model attends over
    n_layers: int = 4
    n_heads: int = 4
    d_model: int = 256           # residual-stream width
    d_ff: int = 1024             # MLP hidden width (4 x d_model)


@dataclass(frozen=True)
class TrainConfig:
    # Data / batching. tokens per step = batch_size * context_length = 16,384.
    batch_size: int = 64
    max_train_tokens: int | None = None  # None = whole train set; set to starve data

    # AdamW optimiser.
    learning_rate: float = 6e-4  # peak LR (reached after warmup)
    min_lr: float = 6e-5         # cosine-decay floor (~10% of peak)
    warmup_iters: int = 200
    max_iters: int = 20_000      # ~327M tokens (~0.6 epoch of the 555M corpus)
    weight_decay: float = 0.1    # applied to matrices only, not norms/biases
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0       # clip global grad norm to this

    # Evaluation, logging, checkpointing (all in optimiser steps).
    eval_interval: int = 500     # how often to measure train/val loss
    eval_iters: int = 100        # batches averaged per loss estimate
    log_interval: int = 20       # how often to print step/loss/lr/speed
    checkpoint_interval: int = 1000

    seed: int = 1337


# Single source of truth other modules import.
model = ModelConfig()
train = TrainConfig()