"""
Nova 1.0 — Humanity's Last Exam (HLE) Automated Evaluation Benchmark Suite
Evaluates Nova 1.0 across 5 core expert academic domains:
1. Advanced Mathematics & Calculus
2. Physics & General Relativity
3. Computer Science & Data Structures
4. Chemistry & Molecular Geometry
5. Formal Logic & Transitive Inferences
"""

import os
import sys
import argparse
import time
import torch

# 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


HLE_BENCHMARK_QUESTIONS = [
    {
        "category": "Mathematics & Calculus",
        "prompt": "Evaluate the definite integral int_0^1 x^7 dx.",
        "expected_keywords": ["1/8", "1 / 8", "0.125", "power rule"],
    },
    {
        "category": "Physics & Quantum Mechanics",
        "prompt": "What is the energy formula of a photon with frequency nu?",
        "expected_keywords": ["h", "nu", "planck", "E = h"],
    },
    {
        "category": "Physics & General Relativity",
        "prompt": "Calculate the Schwarzschild radius of a mass M.",
        "expected_keywords": ["r_s", "2 * G * M", "2GM", "c^2", "c2", "radius"],
    },
    {
        "category": "Computer Science & Algorithms",
        "prompt": "What is the time complexity of building a heap from an array of N elements?",
        "expected_keywords": ["O(N)", "O(n)", "linear", "floyd"],
    },
    {
        "category": "Chemistry & Molecular Geometry",
        "prompt": "What is the molecular geometry and hybridization of sulfur hexafluoride (SF6)?",
        "expected_keywords": ["octahedral", "sp3d2", "sp^3d^2"],
    },
    {
        "category": "Formal Logic & Syllogisms",
        "prompt": "If all A are B, and all B are C, are all A also C?",
        "expected_keywords": ["yes", "transitive", "logically", "follows"],
    },
]


def evaluate_hle_benchmark():
    parser = argparse.ArgumentParser(description="Evaluate Nova 1.0 on HLE Benchmark")
    parser.add_argument("--checkpoint", type=str, default="checkpoints/nova1_final.pt", help="Path to checkpoint")
    parser.add_argument("--tokenizer_path", type=str, default="checkpoints/nova1_gemini_tokenizer.json", help="Tokenizer JSON")
    args = parser.parse_args()

    tokenizer = Nova1Tokenizer.load(args.tokenizer_path)
    checkpoint = torch.load(args.checkpoint, 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(f"\n=======================================================")
    print(f"🏛️ Running Humanity's Last Exam (HLE) Evaluation Benchmark")
    print(f"Model: Nova 1.0 ({sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params) | Device: {config.device.upper()}")
    print(f"=======================================================\n", flush=True)

    passed_count = 0
    total_count = len(HLE_BENCHMARK_QUESTIONS)

    results_table = []

    for idx, test in enumerate(HLE_BENCHMARK_QUESTIONS, 1):
        prompt = test["prompt"]
        category = test["category"]
        expected = test["expected_keywords"]

        start_time = time.time()
        output = generator.generate(
            prompt=prompt,
            max_new_tokens=90,
            temperature=0.4,
            top_k=30,
            repetition_penalty=1.05
        )
        latency = time.time() - start_time

        # Check if output contains expected keywords
        output_lower = output.lower()
        matched = any(kw.lower() in output_lower for kw in expected)
        if matched:
            passed_count += 1
            status = "PASSED ✅"
        else:
            status = "FAILED ❌"

        results_table.append({
            "id": idx,
            "category": category,
            "status": status,
            "latency": latency,
            "output": output
        })

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

    acc_pct = (passed_count / total_count) * 100.0

    print(f"\n=======================================================")
    print(f"📊 HUMANITY'S LAST EXAM (HLE) BENCHMARK SCORECARD")
    print(f"=======================================================")
    print(f"Total Questions Evaluated: {total_count}")
    print(f"Questions Passed:          {passed_count}")
    print(f"Final HLE Benchmark Score: {acc_pct:.1f}%")
    print(f"=======================================================\n", flush=True)

    for r in results_table:
        print(f"\n--- [Question {r['id']}] {r['category']} ({r['status']}) ---")
        print(f"{r['output'].strip()}")
        print("-" * 50)


if __name__ == "__main__":
    evaluate_hle_benchmark()