"""Shared helpers for GPT-U-20M: repo paths, model config and model construction."""
import json
import math
from pathlib import Path

import torch
from transformers import LlamaConfig, LlamaForCausalLM

ROOT = Path(__file__).resolve().parents[1]
CONFIG_PATH = ROOT / "configs" / "model.json"
TOKENIZER_DIR = ROOT / "tokenizer"
DATA_DIR = ROOT / "data" / "tokens"
OUT_DIR = ROOT / "out"
CKPT_DIR = OUT_DIR / "ckpt"
LOG_DIR = ROOT / "logs"

EXPECTED_PARAMS = 20_453_760
INIT_STD = 0.02


def load_config() -> LlamaConfig:
    return LlamaConfig.from_dict(json.loads(CONFIG_PATH.read_text()), attn_implementation="sdpa")


def init_weights(model: LlamaForCausalLM) -> None:
    """N(0, 0.02) everywhere; residual output projections (o_proj, down_proj) scaled by 1/sqrt(2 * n_layers)."""
    resid_std = INIT_STD / math.sqrt(2 * model.config.num_hidden_layers)
    with torch.no_grad():
        for name, module in model.named_modules():
            if name == "lm_head" and model.config.tie_word_embeddings:
                continue  # shares its weight with embed_tokens
            if isinstance(module, torch.nn.Linear):
                module.weight.normal_(0.0, resid_std if name.endswith(("o_proj", "down_proj")) else INIT_STD)
            elif isinstance(module, torch.nn.Embedding):
                module.weight.normal_(0.0, INIT_STD)
            elif type(module).__name__.endswith("RMSNorm"):
                module.weight.fill_(1.0)


def build_model(init: bool = True) -> LlamaForCausalLM:
    # Master weights stay in fp32; bf16 only enters through autocast.
    model = LlamaForCausalLM(load_config()).to(torch.float32)
    if init:
        init_weights(model)
    return model