from ch31_llm_foundations import MiniGPT  # reuse Chapter 31's model


if __name__ == "__main__":
    torch.manual_seed(0)
    vocab_size, block_size = 50, 32
    model = MiniGPT(
        vocab_size=vocab_size, d_model=128, n_heads=4,
        n_layers=4, max_len=block_size,
    )

    total_before = count_total_params(model)
    print(f"Total parameters: {total_before:,}")

    model = apply_lora_to_model(model, r=4, alpha=8, target_attr="qkv")
    trainable = count_trainable_params(model)
    print(f"Trainable parameters after LoRA: {trainable:,}")
    print(f"Fraction trainable: {trainable / total_before:.4%}")

    # verify the frozen base weights genuinely receive no gradient
    x = torch.randint(0, vocab_size, (4, block_size))
    logits = model(x)
    loss = logits.sum()
    loss.backward()

    for module in model.modules():
        if isinstance(module, LoRALinear):
            assert module.base.weight.grad is None
            assert module.A.grad is not None
            assert module.B.grad is not None
    print("Verified: base weights frozen, LoRA params trainable.")
