"""
Post-training validation script.
Tests the trained model by generating text from prompts.

Usage:
    python Base/scripts/validate_model.py --checkpoint_dir Base/out/pretrain/custom-100m-10m-test/final
    python Base/scripts/validate_model.py --checkpoint_dir Base/out/pretrain/custom-100m-3b/final
"""

import argparse
from pathlib import Path

import torch
from litgpt import Tokenizer
from litgpt.config import Config
from litgpt.model import GPT


def load_model(checkpoint_dir: str):
    """Load a pretrained model from a litgpt checkpoint directory."""
    checkpoint_dir = Path(checkpoint_dir)

    if not checkpoint_dir.exists():
        raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")

    config = Config.from_checkpoint(checkpoint_dir)
    model_path = checkpoint_dir / "lit_model.pth"

    if not model_path.exists():
        raise FileNotFoundError(f"Model file not found: {model_path}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    with torch.device("meta"):
        model = GPT(config)

    checkpoint = torch.load(str(model_path), map_location="cpu", weights_only=False)

    # Handle both direct state_dict and wrapped checkpoint formats
    if "model" in checkpoint:
        state_dict = checkpoint["model"]
    else:
        state_dict = checkpoint

    model.load_state_dict(state_dict, assign=True)
    model = model.to(device)
    model.eval()

    return model, config, device


def generate_text(model, tokenizer, device, prompt, max_new_tokens=200, temperature=0.8, top_k=50):
    """Generate text from a prompt."""
    input_ids = tokenizer.encode(prompt, device=device).unsqueeze(0)

    with torch.no_grad():
        for _ in range(max_new_tokens):
            # Crop to block_size if needed
            idx_cond = input_ids[:, -model.max_seq_length:]
            logits = model(idx_cond)
            logits = logits[:, -1, :] / temperature

            # Top-k filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")

            probs = torch.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_id], dim=1)

            # Stop on EOS
            if next_id.item() == tokenizer.eos_id:
                break

    return tokenizer.decode(input_ids[0])


def main():
    parser = argparse.ArgumentParser(description="Validate a pretrained LitGPT model")
    parser.add_argument(
        "--checkpoint_dir",
        type=str,
        required=True,
        help="Path to the checkpoint directory (e.g., Base/out/pretrain/custom-100m-10m-test/final)",
    )
    parser.add_argument(
        "--tokenizer_dir",
        type=str,
        default="Base/checkpoints/EleutherAI/pythia-160m",
        help="Path to the tokenizer directory",
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=200,
        help="Maximum tokens to generate",
    )
    args = parser.parse_args()

    print(f"Loading model from: {args.checkpoint_dir}")
    model, config, device = load_model(args.checkpoint_dir)

    print(f"\nModel config: {config.name}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
    print(f"Block size: {config.block_size}")

    tokenizer = Tokenizer(Path(args.tokenizer_dir))

    test_prompts = [
        "The future of artificial intelligence",
        "In a groundbreaking study, researchers found that",
        "The most important thing about education is",
        "Once upon a time, in a land far away,",
    ]

    print(f"\n{'='*60}")
    print("GENERATION TESTS")
    print(f"{'='*60}")

    for i, prompt in enumerate(test_prompts, 1):
        print(f"\n--- Prompt {i}: \"{prompt}\" ---")
        output = generate_text(
            model, tokenizer, device, prompt,
            max_new_tokens=args.max_new_tokens,
        )
        print(f"Generated:\n{output}\n")

    print(f"{'='*60}")
    print("Validation complete!")
    print("Note: A freshly pretrained model will produce semi-coherent text.")
    print("Quality improves with more training data and compute.")


if __name__ == "__main__":
    main()