"""
Eve-2-MoE Inference
===================
Quick generation script. Works with local weights or HuggingFace download.

Usage:
    python generate.py --prompt "The future of AI is"
    python generate.py --prompt "The future of AI is" --model_path ./model_final/pytorch_model.bin
    python generate.py --prompt "The future of AI is" --hf_repo anthonym21/Eve-2-MoE-250M
"""

import argparse
import torch
import tiktoken
from modeling_eve import ModelConfig, DeepSeekMoE


def load_model(model_path: str = None, hf_repo: str = None, device: str = "cuda"):
    config = ModelConfig()
    model = DeepSeekMoE(config)

    if hf_repo:
        from huggingface_hub import hf_hub_download
        model_path = hf_hub_download(repo_id=hf_repo, filename="pytorch_model.bin")

    if model_path:
        state_dict = torch.load(model_path, map_location=device, weights_only=True)
        model.load_state_dict(state_dict)

    return model.to(device).eval()


def generate_streaming(model, prompt: str, max_tokens: int = 200,
                       temperature: float = 0.8, top_k: int = 50, device: str = "cuda"):
    enc = tiktoken.get_encoding("gpt2")
    tokens = torch.tensor(enc.encode(prompt), dtype=torch.long, device=device).unsqueeze(0)

    print(prompt, end="", flush=True)

    with torch.no_grad():
        for _ in range(max_tokens):
            idx_cond = tokens[:, -model.config.block_size:]

            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=(device == "cuda")):
                logits, _ = model(idx_cond)

            logits = logits[:, -1, :] / temperature

            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)
            idx_next = torch.multinomial(probs, num_samples=1)
            tokens = torch.cat((tokens, idx_next), dim=1)

            print(enc.decode([idx_next.item()]), end="", flush=True)

    print("\n")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--prompt", type=str, default="The future of artificial intelligence is")
    p.add_argument("--model_path", type=str, default=None)
    p.add_argument("--hf_repo", type=str, default=None)
    p.add_argument("--max_tokens", type=int, default=200)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top_k", type=int, default=50)
    p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
    args = p.parse_args()

    if not args.model_path and not args.hf_repo:
        args.hf_repo = "anthonym21/Eve-2-MoE-250M"

    print(f"Loading model on {args.device}...")
    model = load_model(args.model_path, args.hf_repo, args.device)
    param_count = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {param_count / 1e6:.2f}M\n")

    generate_streaming(model, args.prompt, args.max_tokens, args.temperature, args.top_k, args.device)


if __name__ == "__main__":
    main()