"""
Self-contained loader for this checkpoint.

Reconstructs the model architecture and tokenizer entirely from
config.json (no hardcoded hyperparameters here), loads the trained
weights, and generates a sample -- proving this folder is sufficient on
its own, with no dependency on the original training repo.

Usage:
    python load_model.py
"""

from __future__ import annotations

import json
from pathlib import Path

import torch

from bigram import BigramLanguageModel

HERE = Path(__file__).resolve().parent


def _load_weights(path: Path) -> dict[str, torch.Tensor]:
    if path.suffix == ".safetensors":
        from safetensors.torch import load_file

        return load_file(path)
    return torch.load(path, map_location="cpu", weights_only=True)


def load_model_and_tokenizer(dir_path: Path = HERE):
    config = json.loads((dir_path / "config.json").read_text())

    if config["tokenizer_type"] == "char":
        from tokenizer import CharTokenizer as TokenizerClass
    elif config["tokenizer_type"] == "bpe":
        from bpe_tokenizer import BPETokenizer as TokenizerClass
    else:
        raise ValueError(f"Unknown tokenizer_type: {config['tokenizer_type']!r}")
    tokenizer = TokenizerClass.load(dir_path / config["tokenizer_file"])

    model = BigramLanguageModel(
        vocab_size=config["vocab_size"],
        block_size=config["block_size"],
        n_embd=config["n_embd"],
        n_head=config["n_head"],
        n_layer=config["n_layer"],
    )
    state_dict = _load_weights(dir_path / config["weights_file"])
    model.load_state_dict(state_dict)
    model.eval()

    return model, tokenizer, config


def main() -> None:
    model, tokenizer, config = load_model_and_tokenizer()

    n_params = sum(p.numel() for p in model.parameters())
    print(f"Loaded model: {n_params:,} parameters (config.json says {config['parameter_count']:,})")
    assert n_params == config["parameter_count"], "Reconstructed model doesn't match config!"

    context = torch.zeros((1, 1), dtype=torch.long)
    print("\n----- generated sample -----")
    sample = tokenizer.decode(model.generate(context, max_new_tokens=400)[0].tolist())
    print(sample)


if __name__ == "__main__":
    main()