"""
model_factory.py — model config builder.

The old presets dict hardcoded vocab_size and max_seq_len per architecture.
That coupled the model to a specific tokenizer and sequence length, which is
exactly the thing we're moving away from now that vocab is selectable
(32k vs 64k) and T may grow.

New API: build_model_config(model_type, size, vocab_size, max_seq_len) returns
a ModelArgs with the architectural shape from the preset table and vocab_size
+ max_seq_len from the caller.
"""

from llama3 import ModelArgs


# ----------------------------------------------------------------------------
#  Architectural presets — shape only. No vocab_size, no max_seq_len.
# ----------------------------------------------------------------------------
#
# Each entry holds the model-shape hyperparameters that don't depend on the
# tokenizer or context length: layer count, head count, hidden dim, MLP
# intermediate size, and KV head count.
#
# The "llama-iwslt" entries match the previous model_presets sizes, just with
# vocab_size factored out. The "llama" entries (pre-IWSLT vocab 50304) are
# kept for backward compat with old checkpoints but new training should use
# llama-iwslt.

ARCH_PRESETS: dict[str, dict[str, dict]] = {
    "llama-iwslt": {
        "124m": dict(n_layers=12, n_heads=12, dim=768,  intermediate_size=4 * 768,  n_kv_heads=12),
        "500m": dict(n_layers=32, n_heads=16, dim=1024, intermediate_size=3072,     n_kv_heads=4),
        "978m": dict(n_layers=36, n_heads=20, dim=1280, intermediate_size=5120,     n_kv_heads=4),
        "1b":   dict(n_layers=36, n_heads=20, dim=1280, intermediate_size=5120,     n_kv_heads=4),
        "2b":   dict(n_layers=36, n_heads=32, dim=2048, intermediate_size=5120,     n_kv_heads=8),
    },
}


def build_model_config(
    model_type: str,
    size: str,
    vocab_size: int,
    max_seq_len: int,
) -> ModelArgs:
    """Build a ModelArgs combining architectural preset + runtime vocab/seq_len."""
    if model_type not in ARCH_PRESETS:
        raise ValueError(f"unknown model_type {model_type!r}; "
                         f"expected one of {list(ARCH_PRESETS)}")
    if size not in ARCH_PRESETS[model_type]:
        raise ValueError(f"unknown size {size!r} for {model_type}; "
                         f"expected one of {list(ARCH_PRESETS[model_type])}")
    arch = ARCH_PRESETS[model_type][size]
    return ModelArgs(
        vocab_size=vocab_size,
        max_seq_len=max_seq_len,
        **arch,
    )


# ----------------------------------------------------------------------------
#  Checkpoint loader (kept for backward compat with utils.completion etc.)
# ----------------------------------------------------------------------------

def load_model(model, path, load_dict_only=False):
    """Load a model checkpoint into an already-constructed model.

    Handles both:
      - new checkpoints saved by the new train.py (keys = state_dict layout)
      - old checkpoints with torch.compile's '_orig_mod.' prefix
    """
    import torch
    ckpt = torch.load(path, map_location=torch.device("cpu"), weights_only=False)
    sd = ckpt["model"] if "model" in ckpt else ckpt
    # Strip torch.compile prefix
    sd = {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
    # Strip DDP prefix
    sd = {k.replace("module.", ""): v for k, v in sd.items()}
    model.load_state_dict(sd)
    step = ckpt.get("step", 0) if isinstance(ckpt, dict) else 0
    if not load_dict_only:
        model.train()
    return model, step