import os
import shutil

import torch
from transformers import AutoModelForCausalLM

from configuration_minigpt import MiniGPTConfig
from modeling_minigpt import MiniGPTForCausalLM

OUT_DIR = "runs/check_model"
EXPECTED_PARAMS = 31_433_920


def param_breakdown(config):
    embed = config.vocab_size * config.d_model
    qkvo = 4 * config.d_model * config.d_model
    mlp = 2 * config.d_model * config.d_ff
    norms = (2 * config.n_layer + 1) * config.d_model
    return embed, qkvo + mlp, norms


def main():
    torch.manual_seed(0)
    config = MiniGPTConfig()
    model = MiniGPTForCausalLM(config)
    model.eval()

    embed, block, norms = param_breakdown(config)
    n_params = model.num_parameters()
    print(f"embedding      {embed:>10,}")
    print(f"blocks         {block:>10,}")
    print(f"norms          {norms:>10,}")
    print(f"total          {n_params:>10,}")
    assert n_params == EXPECTED_PARAMS, f"{n_params} != {EXPECTED_PARAMS}"

    assert model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr()

    input_ids = torch.randint(0, config.vocab_size, (2, 32))
    with torch.no_grad():
        logits = model(input_ids).logits
    assert logits.shape == (2, 32, config.vocab_size), logits.shape

    edited = input_ids.clone()
    edited[:, 17] = (edited[:, 17] + 1) % config.vocab_size
    with torch.no_grad():
        logits_edited = model(edited).logits
    assert torch.equal(logits[:, :17], logits_edited[:, :17])
    assert not torch.equal(logits[:, 17], logits_edited[:, 17])

    if os.path.isdir(OUT_DIR):
        shutil.rmtree(OUT_DIR)
    model.save_pretrained(OUT_DIR)

    for name in [
        "config.json",
        "model.safetensors",
        "configuration_minigpt.py",
        "modeling_minigpt.py",
    ]:
        path = os.path.join(OUT_DIR, name)
        assert os.path.isfile(path), f"missing {name}"

    reloaded = AutoModelForCausalLM.from_pretrained(OUT_DIR, trust_remote_code=True)
    reloaded.eval()
    with torch.no_grad():
        reloaded_logits = reloaded(input_ids).logits
    assert torch.allclose(logits, reloaded_logits, atol=1e-5), "logits differ after reload"

    print("PASS")


if __name__ == "__main__":
    main()