"""
Nova 1.0 — Official Humanity's Last Exam (cais/hle) Benchmark Evaluation Script
Evaluates Nova 1.0 across official questions from the cais/hle Hugging Face dataset.
"""

import os
import sys
import argparse
import time
import torch
from datasets import load_dataset

# Add project root to sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from config.model_config import Nova1Config
from src.tokenizer.bpe_tokenizer import Nova1Tokenizer
from src.model.nova1_hrm import Nova1HRM
from src.inference.generate import Nova1Generator


def run_cais_hle_evaluation(max_eval_samples: int = 50):
    token = os.environ.get("HF_TOKEN")
    tokenizer_path = "checkpoints/nova1_gemini_tokenizer.json"
    checkpoint_path = "checkpoints/nova1_final.pt"

    tokenizer = Nova1Tokenizer.load(tokenizer_path)
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    config: Nova1Config = checkpoint.get("config", Nova1Config(vocab_size=tokenizer.vocab_size))
    config.device = "cuda" if torch.cuda.is_available() else "cpu"
    config.dtype = "bfloat16"

    model = Nova1HRM(config)
    state_dict = checkpoint["model_state"] if "model_state" in checkpoint else checkpoint["model_state_dict"]
    model.load_state_dict(state_dict)
    model.to(config.device)
    model.eval()

    generator = Nova1Generator(model, tokenizer, config)

    print("Loading official 'cais/hle' dataset from Hugging Face Hub...", flush=True)
    hle_ds = load_dataset("cais/hle", split="test", token=token)
    print(f"Loaded cais/hle test set ({len(hle_ds):,} total samples).", flush=True)

    text_samples = [item for item in hle_ds if item.get("image") is None][:max_eval_samples]
    print(f"Evaluating Nova 1.0 across {len(text_samples)} textual HLE questions...", flush=True)

    passed_count = 0
    category_scores = {}

    for idx, item in enumerate(text_samples, 1):
        q_text = item.get("question", "")
        answer = str(item.get("answer", "")).strip()
        cat = item.get("category", "General STEM")

        if cat not in category_scores:
            category_scores[cat] = {"passed": 0, "total": 0}
        category_scores[cat]["total"] += 1

        start_t = time.time()
        output = generator.generate(
            prompt=q_text,
            max_new_tokens=80,
            temperature=0.4,
            top_k=30,
            repetition_penalty=1.05
        )
        latency = time.time() - start_t

        output_lower = output.lower()
        ans_lower = answer.lower()

        # Check if answer or key terms match
        is_pass = (ans_lower in output_lower) or any(tok in output_lower for tok in ans_lower.split() if len(tok) > 3)
        if is_pass:
            passed_count += 1
            category_scores[cat]["passed"] += 1
            status = "PASSED ✅"
        else:
            status = "FAILED ❌"

        print(f"[{idx}/{len(text_samples)}] Category: {cat} -> {status} ({latency:.2f}s)", flush=True)

    total_pct = (passed_count / len(text_samples)) * 100.0 if text_samples else 0.0
    print(f"\n=======================================================")
    print(f"🏆 OFFICIAL CAIS/HLE BENCHMARK ACCURACY: {total_pct:.1f}%")
    print(f"=======================================================\n", flush=True)

    return total_pct, category_scores


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate Nova 1.0 on cais/hle")
    parser.add_argument("--samples", type=int, default=30, help="Number of samples to evaluate")
    args = parser.parse_args()

    run_cais_hle_evaluation(args.samples)