#!/usr/bin/env python3

import torch
import torch.nn.functional as F
import argparse
import numpy as np
from pathlib import Path
from data import Dataset


def load_latest_model(model_dir="."):
    """Load the most recent model checkpoint."""
    model_files = list(Path(model_dir).glob("model_epoch_*.pth"))
    if not model_files:
        raise FileNotFoundError("No model checkpoints found!")

    # Sort by modification time to get the latest
    latest_model = max(model_files, key=lambda x: x.stat().st_mtime)
    print(f"Loading model: {latest_model}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(latest_model, map_location=device, weights_only=False)
    model.eval()
    return model, device


def generate_text(
    model, dataset, device, max_length=200, temperature=1.0, seed_text=""
):
    """Generate text character by character."""
    model.eval()

    # Convert seed text to indices
    if seed_text:
        # Create vocab mapping
        vocab_to_idx = {char: idx for idx, char in enumerate(dataset.vocab)}
        try:
            seed_indices = [vocab_to_idx[char] for char in seed_text]
        except KeyError as e:
            print(f"Warning: Character {e} not in vocabulary, using empty seed")
            seed_indices = []
    else:
        seed_indices = []

    # Start with seed or random character
    if seed_indices:
        generated = seed_indices.copy()
    else:
        # Start with a random character from vocabulary
        generated = [np.random.randint(0, len(dataset.vocab))]

    with torch.no_grad():
        for _ in range(max_length - len(generated)):
            # Convert to tensor and get the last block_size tokens
            context = generated[-model.block_size :]
            context_tensor = torch.tensor(
                context, dtype=torch.long, device=device
            ).unsqueeze(0)

            # Get predictions
            logits = model(context_tensor)

            # Focus on the last token's predictions
            logits = logits[0, -1, :] / temperature

            # Apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1)

            # Sample from the distribution
            next_token = torch.multinomial(probs, num_samples=1).item()
            generated.append(next_token)

            # Stop if we generate too much or hit some limit
            if len(generated) >= max_length:
                break

    # Convert indices back to text
    generated_text = "".join([dataset.vocab[idx] for idx in generated])
    return generated_text


def main():
    parser = argparse.ArgumentParser(
        description="Generate text using trained language model"
    )
    parser.add_argument("N", type=int, help="Number of text samples to generate")
    parser.add_argument(
        "--length", type=int, default=200, help="Length of each generated sample"
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=1.0,
        help="Sampling temperature (higher = more random)",
    )
    parser.add_argument(
        "--seed", type=str, default="", help="Seed text to start generation"
    )

    args = parser.parse_args()

    print(f"Generating {args.N} samples...")

    # Load dataset to get vocabulary
    print("Loading dataset...")
    dataset = Dataset()
    print(f"Vocabulary size: {len(dataset.vocab)}")

    # Load the latest model
    model, device = load_latest_model()
    print(f"Using device: {device}")

    # Generate N samples
    for i in range(args.N):
        print(f"\n{'='*50}")
        print(f"Generated Sample {i+1}:")
        print(f"{'='*50}")

        generated_text = generate_text(
            model,
            dataset,
            device,
            max_length=args.length,
            temperature=args.temperature,
            seed_text=args.seed,
        )

        print(generated_text)
        print(f"{'='*50}")


if __name__ == "__main__":
    main()
