#!/usr/bin/env python3
"""
MarsNext Phobos-1 — Head-to-Head Multi-Model Benchmark Harness
(`phobos-1/models/phobos_scratch_v1/run_multi_model_comparison.py`)
===================================================================
Benchmarks `MarsNext/Phobos-1` (`52.22M` Active / `96.26M` Total Params, `1.56B` tokens)
head-to-head against established open-source models in the `70M - 165M` active-parameter class:
  1. `EleutherAI/pythia-70m-deduped`  (`70.4M` active params, `300B` tokens)
  2. `openai-community/gpt2`          (`124.4M` active params, `~40B` tokens)
  3. `EleutherAI/pythia-160m-deduped` (`162.3M` active params, `300B` tokens)
  4. `HuggingFaceTB/SmolLM2-135M`     (`134.5M` active params, `2,000B` tokens)

Evaluates every model on:
  A. Standardized 4-Choice Log-Likelihood Benchmark (`ARC-Easy / SciQ / PIQA / CS-Math`, 20 items)
     using tokenizer-independent length-normalized log-probability.
  B. 25-Task Generative Suite (`Python Pass@1 exec()`, `SQL`, `STEM`, `Math`, `<think>` Reasoning).
  C. Parameter Efficiency (`Score per 100M Active Params`) and Apple Silicon (`MPS`) Speed (`tok/s`).

Strictly <= 290 lines (`RULE[user_global]`).
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path
from typing import Any, Callable, Dict, List

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

_THIS_DIR = Path(__file__).resolve().parent
if str(_THIS_DIR) not in sys.path:
    sys.path.insert(0, str(_THIS_DIR))

try:
    from chat_cli import load_phobos_engine, resolve_default_tokenizer  # type: ignore
except ImportError:
    load_phobos_engine = None  # type: ignore
    resolve_default_tokenizer = None  # type: ignore
from run_phobos1_benchmark_suite import (  # type: ignore
    BASE_CKPT,
    FINAL_CKPT,
    MATH_TASKS,
    OUT_JSON as PHOBOS_BENCH_JSON,
    PYTHON_TASKS,
    SQL_TASKS,
    STEM_TASKS,
    THINK_TASKS,
    _check_required_any,
    _test_python_exec,
)

MULTI_OUT_JSON = _THIS_DIR / "weights" / "multi_model_head_to_head_benchmark.json"

# 20 Standardized 4-Choice Questions (ARC-Easy, SciQ, PIQA, CS & Math Fundamentals)
MC_QUESTIONS: List[Dict[str, Any]] = [
    {"id": "ARC-01", "q": "Which organelle is the primary site of ATP production in eukaryotic muscle cells?", "choices": ["Golgi apparatus", "Mitochondrion", "Lysosome", "Ribosome"], "ans": 1},
    {"id": "ARC-02", "q": "What force keeps the moons Phobos and Deimos in orbit around the planet Mars?", "choices": ["Magnetic repulsion", "Solar wind", "Gravity", "Atmospheric friction"], "ans": 2},
    {"id": "ARC-03", "q": "When hydrochloric acid (HCl) reacts with sodium hydroxide (NaOH), the products are water and:", "choices": ["Sodium chloride", "Sulfuric acid", "Ammonia", "Carbon dioxide"], "ans": 0},
    {"id": "ARC-04", "q": "In plants, chlorophyll absorbs light energy primarily to drive which biological process?", "choices": ["Fermentation", "Photosynthesis", "Mitosis", "Transpiration"], "ans": 1},
    {"id": "ARC-05", "q": "Which subatomic particle carries a negative electric charge?", "choices": ["Proton", "Neutron", "Electron", "Photon"], "ans": 2},
    {"id": "SCIQ-01", "q": "What is the SI unit of electrical resistance?", "choices": ["Volt", "Ampere", "Ohm", "Watt"], "ans": 2},
    {"id": "SCIQ-02", "q": "During DNA replication, which enzyme synthesizes the new complementary DNA strand?", "choices": ["RNAse", "DNA polymerase", "Lipase", "Amylase"], "ans": 1},
    {"id": "SCIQ-03", "q": "Which gas makes up the largest percentage of Mars's thin atmosphere?", "choices": ["Oxygen", "Nitrogen", "Carbon dioxide", "Hydrogen"], "ans": 2},
    {"id": "SCIQ-04", "q": "If an object has a mass of 5 kg and accelerates at 4 m/s^2, what is the net force acting on it?", "choices": ["9 N", "20 N", "1.25 N", "40 N"], "ans": 1},
    {"id": "SCIQ-05", "q": "At standard atmospheric pressure, pure water boils at what temperature on the Celsius scale?", "choices": ["0 degrees Celsius", "50 degrees Celsius", "100 degrees Celsius", "212 degrees Celsius"], "ans": 2},
    {"id": "PIQA-01", "q": "To prevent an iron nail from rusting outdoors, what should you coat it with?", "choices": ["Salt water", "Paint or oil", "Vinegar", "Wet sand"], "ans": 1},
    {"id": "PIQA-02", "q": "If you want to cool down a cup of hot tea faster, you should:", "choices": ["Put a tight lid on the cup", "Stir it with a metal spoon and blow across the surface", "Wrap the cup in wool", "Place it inside a closed microwave"], "ans": 1},
    {"id": "PIQA-03", "q": "To loosen a tightly stuck metal jar lid, it helps to:", "choices": ["Freeze the entire jar", "Run warm water over the metal lid so it expands", "Hammer the glass bottom", "Coat the lid with butter"], "ans": 1},
    {"id": "CS-01", "q": "In Python, which operator returns the remainder of integer division?", "choices": ["//", "%", "**", "&"], "ans": 1},
    {"id": "CS-02", "q": "What is the worst-case time complexity of binary search on a sorted array of N elements?", "choices": ["O(N^2)", "O(N)", "O(log N)", "O(N!)" ], "ans": 2},
    {"id": "CS-03", "q": "In SQL, which clause is used to filter rows before grouping?", "choices": ["ORDER BY", "WHERE", "GROUP BY", "LIMIT"], "ans": 1},
    {"id": "CS-04", "q": "Which data structure operates on a Last-In, First-Out (LIFO) principle?", "choices": ["Queue", "Stack", "Hash Table", "Binary Search Tree"], "ans": 1},
    {"id": "MATH-MC1", "q": "If 3x + 7 = 22, what is the value of x?", "choices": ["3", "5", "7", "15"], "ans": 1},
    {"id": "MATH-MC2", "q": "According to the calculus power rule, the first derivative of f(x) = x^3 is:", "choices": ["x^2", "3x^2", "3x^3", "x^4 / 4"], "ans": 1},
    {"id": "MATH-MC3", "q": "A right triangle has legs of length 3 and 4. What is the length of its hypotenuse?", "choices": ["5", "6", "7", "12"], "ans": 0},
]

EXTERNAL_MODELS: List[Dict[str, Any]] = [
    {"id": "EleutherAI/pythia-70m-deduped", "org": "EleutherAI", "train_tokens": "300B"},
    {"id": "openai-community/gpt2", "org": "OpenAI", "train_tokens": "40B"},
    {"id": "EleutherAI/pythia-160m-deduped", "org": "EleutherAI", "train_tokens": "300B"},
    {"id": "HuggingFaceTB/SmolLM2-135M", "org": "HuggingFace", "train_tokens": "2,000B"},
]


@torch.no_grad()
def score_mc_choice(
    forward_logits_fn: Callable[[torch.Tensor], torch.Tensor],
    encode_fn: Callable[[str], List[int]],
    device: torch.device,
    question: str,
    choice: str,
) -> float:
    """Computes length-normalized log-likelihood of `choice` given `question`."""
    prefix = f"Question: {question}\nAnswer:"
    full = f"Question: {question}\nAnswer: {choice}"
    p_ids = encode_fn(prefix)
    f_ids = encode_fn(full)
    if len(f_ids) <= len(p_ids):
        f_ids = p_ids + encode_fn(" " + choice)
    c_len = max(1, len(f_ids) - len(p_ids))
    inp = torch.tensor([f_ids], dtype=torch.long, device=device)
    logits = forward_logits_fn(inp)[0].float()
    log_probs = F.log_softmax(logits, dim=-1)
    start_idx = len(f_ids) - c_len
    total_lp = 0.0
    for pos in range(start_idx, len(f_ids)):
        if pos - 1 >= 0:
            total_lp += float(log_probs[pos - 1, f_ids[pos]].item())
    return total_lp / c_len


@torch.no_grad()
def evaluate_mc_benchmark(
    forward_logits_fn: Callable[[torch.Tensor], torch.Tensor],
    encode_fn: Callable[[str], List[int]],
    device: torch.device,
) -> Dict[str, Any]:
    """Evaluates all 20 multiple-choice items via length-normalized log-likelihood."""
    correct = 0
    details = []
    for item in MC_QUESTIONS:
        scores = [
            score_mc_choice(forward_logits_fn, encode_fn, device, item["q"], ch)
            for ch in item["choices"]
        ]
        pred = int(max(range(len(scores)), key=lambda i: scores[i]))
        hit = pred == item["ans"]
        correct += int(hit)
        details.append({"id": item["id"], "pred": pred, "ans": item["ans"], "correct": hit})
    return {"correct": correct, "total": len(MC_QUESTIONS), "acc_pct": round(100.0 * correct / len(MC_QUESTIONS), 1), "details": details}


@torch.no_grad()
def generate_hf_reply(
    model: Any,
    tok: Any,
    device: torch.device,
    prompt_text: str,
    force_think: bool = False,
    max_new_tokens: int = 96,
) -> tuple[str, int, float]:
    """Greedy generation with repetition penalty (1.15) for standard HuggingFace CausalLMs."""
    prefix = "<think>\n" if force_think else ""
    formatted = f"User: {prompt_text.strip()}\nAssistant:\n{prefix}"
    input_ids = tok.encode(formatted, add_special_tokens=False)
    ids = torch.tensor([input_ids], dtype=torch.long, device=device)
    gen_ids: List[int] = tok.encode(prefix, add_special_tokens=False) if force_think else []
    eos_id = tok.eos_token_id
    t0 = time.time()
    for _ in range(max_new_tokens):
        logits = model(ids[:, -1024:]).logits[0, -1, :].float()
        for seen in set(ids[0].tolist()[-64:]):
            logits[seen] = logits[seen] / 1.15 if logits[seen] > 0 else logits[seen] * 1.15
        nxt = int(torch.argmax(logits).item())
        if eos_id is not None and nxt == eos_id:
            break
        gen_ids.append(nxt)
        ids = torch.cat([ids, torch.tensor([[nxt]], dtype=torch.long, device=device)], dim=1)
    el = max(1e-6, time.time() - t0)
    return tok.decode(gen_ids, skip_special_tokens=False), len(gen_ids), el


def evaluate_external_hf_model(spec: Dict[str, Any], device: torch.device) -> Dict[str, Any]:
    """Downloads/loads an external HuggingFace model and runs both MC-20 and Generative-25 suites."""
    model_id = spec["id"]
    print(f"=== Loading External Competitor: {model_id} ===", flush=True)
    tok = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32).to(device).eval()
    n_params = sum(p.numel() for p in model.parameters())

    mc_res = evaluate_mc_benchmark(
        lambda x: model(x).logits,
        lambda s: tok.encode(s, add_special_tokens=False),
        device,
    )

    total_tok = 0
    total_sec = 0.0
    py_pass = 0
    for t in PYTHON_TASKS:
        rep, n_t, el = generate_hf_reply(model, tok, device, t["prompt"], False, 96)
        total_tok += n_t
        total_sec += el
        py_pass += int(_test_python_exec(rep, t["check"]))

    sql_pass = 0
    for t in SQL_TASKS:
        rep, n_t, el = generate_hf_reply(model, tok, device, t["prompt"], False, 72)
        total_tok += n_t
        total_sec += el
        sql_pass += int(all(req in rep.lower() for req in t["required_tokens"]))

    stem_pass = 0
    for t in STEM_TASKS:
        rep, n_t, el = generate_hf_reply(model, tok, device, t["prompt"], False, 96)
        total_tok += n_t
        total_sec += el
        stem_pass += int(_check_required_any(rep, t["required_any"]))

    math_pass = 0
    for t in MATH_TASKS:
        rep, n_t, el = generate_hf_reply(model, tok, device, t["prompt"], False, 96)
        total_tok += n_t
        total_sec += el
        math_pass += int(_check_required_any(rep, t["required_any"]))

    think_pass = 0
    for t in THINK_TASKS:
        rep, n_t, el = generate_hf_reply(model, tok, device, t["prompt"], True, 96)
        total_tok += n_t
        total_sec += el
        think_pass += int(_check_required_any(rep, t["required_any"]))

    gen_total = py_pass + sql_pass + stem_pass + math_pass + think_pass
    gen_pct = round(100.0 * gen_total / 25.0, 1)
    combined_pct = round((mc_res["acc_pct"] + gen_pct) / 2.0, 1)
    speed = round(total_tok / max(1e-6, total_sec), 1)

    res = {
        "model_id": model_id,
        "org": spec["org"],
        "total_params_m": round(n_params / 1e6, 2),
        "active_params_m": round(n_params / 1e6, 2),
        "train_tokens": spec["train_tokens"],
        "mc20_acc_pct": mc_res["acc_pct"],
        "mc20_correct": f"{mc_res['correct']}/20",
        "gen25_score_pct": gen_pct,
        "gen25_passed": f"{gen_total}/25",
        "python_pass_at_1": f"{py_pass}/5 ({py_pass*20}%)",
        "sql_synthesis": f"{sql_pass}/5 ({sql_pass*20}%)",
        "stem_qa": f"{stem_pass}/5 ({stem_pass*20}%)",
        "math_symbolic": f"{math_pass}/5 ({math_pass*20}%)",
        "think_planning": f"{think_pass}/5 ({think_pass*20}%)",
        "combined_avg_pct": combined_pct,
        "speed_tok_per_s": speed,
    }
    print(f"[{model_id}] MC-20={mc_res['acc_pct']}% | Gen-25={gen_pct}% ({gen_total}/25) | Combined={combined_pct}% | {speed} tok/s", flush=True)
    del model
    if torch.backends.mps.is_available():
        torch.mps.empty_cache()
    return res


def main() -> None:
    t0 = time.time()
    if "--colab-ext" in sys.argv:
        dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("Colab GPU:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu", flush=True)
        ext_res = [evaluate_external_hf_model(spec, dev) for spec in EXTERNAL_MODELS]
        Path("/content/colab_ext_bench.json").write_text(json.dumps(ext_res, indent=2), encoding="utf-8")
        print("=== COLAB EXTERNAL BENCHMARK SAVED ===", flush=True)
        return

    tok_path = resolve_default_tokenizer()
    phobos_bench = json.loads(PHOBOS_BENCH_JSON.read_text())

    # Evaluate Phobos-1 Base and Phobos-1 Final on MC-20
    m_base, tok_p, dev = load_phobos_engine(BASE_CKPT, tok_path)
    mc_base = evaluate_mc_benchmark(
        lambda x: m_base(x)["logits"],
        lambda s: tok_p.encode(s, max_len=512),
        dev,
    )
    del m_base

    m_final, tok_p, dev = load_phobos_engine(FINAL_CKPT, tok_path)
    mc_final = evaluate_mc_benchmark(
        lambda x: m_final(x)["logits"],
        lambda s: tok_p.encode(s, max_len=512),
        dev,
    )
    del m_final

    b_suite = phobos_bench["base_step33977_suite"]
    f_suite = phobos_bench["phobos1_final_suite"]

    leaderboard: List[Dict[str, Any]] = [
        {
            "model_id": "marsnext/Phobos-1 (Final Post-Trained)",
            "org": "MarsNext",
            "total_params_m": 96.26,
            "active_params_m": 52.22,
            "train_tokens": "1.56B",
            "mc20_acc_pct": mc_final["acc_pct"],
            "mc20_correct": f"{mc_final['correct']}/20",
            "gen25_score_pct": f_suite["overall_score_pct"],
            "gen25_passed": f"{f_suite['overall_passed']}/25",
            "python_pass_at_1": f"{f_suite['suites']['python_pass_at_1']['passed']}/5 ({int(f_suite['suites']['python_pass_at_1']['score_pct'])}%)",
            "sql_synthesis": f"{f_suite['suites']['sql_query_synthesis']['passed']}/5 ({int(f_suite['suites']['sql_query_synthesis']['score_pct'])}%)",
            "stem_qa": f"{f_suite['suites']['stem_science_qa']['passed']}/5 ({int(f_suite['suites']['stem_science_qa']['score_pct'])}%)",
            "math_symbolic": f"{f_suite['suites']['math_symbolic_reasoning']['passed']}/5 ({int(f_suite['suites']['math_symbolic_reasoning']['score_pct'])}%)",
            "think_planning": f"{f_suite['suites']['system2_think_planning']['passed']}/5 ({int(f_suite['suites']['system2_think_planning']['score_pct'])}%)",
            "combined_avg_pct": round((mc_final["acc_pct"] + f_suite["overall_score_pct"]) / 2.0, 1),
            "speed_tok_per_s": f_suite["avg_tok_per_sec"],
        },
        {
            "model_id": "marsnext/Phobos-1 (Base Step 33,977)",
            "org": "MarsNext",
            "total_params_m": 96.26,
            "active_params_m": 52.22,
            "train_tokens": "1.56B",
            "mc20_acc_pct": mc_base["acc_pct"],
            "mc20_correct": f"{mc_base['correct']}/20",
            "gen25_score_pct": b_suite["overall_score_pct"],
            "gen25_passed": f"{b_suite['overall_passed']}/25",
            "python_pass_at_1": f"{b_suite['suites']['python_pass_at_1']['passed']}/5 ({int(b_suite['suites']['python_pass_at_1']['score_pct'])}%)",
            "sql_synthesis": f"{b_suite['suites']['sql_query_synthesis']['passed']}/5 ({int(b_suite['suites']['sql_query_synthesis']['score_pct'])}%)",
            "stem_qa": f"{b_suite['suites']['stem_science_qa']['passed']}/5 ({int(b_suite['suites']['stem_science_qa']['score_pct'])}%)",
            "math_symbolic": f"{b_suite['suites']['math_symbolic_reasoning']['passed']}/5 ({int(b_suite['suites']['math_symbolic_reasoning']['score_pct'])}%)",
            "think_planning": f"{b_suite['suites']['system2_think_planning']['passed']}/5 ({int(b_suite['suites']['system2_think_planning']['score_pct'])}%)",
            "combined_avg_pct": round((mc_base["acc_pct"] + b_suite["overall_score_pct"]) / 2.0, 1),
            "speed_tok_per_s": b_suite["avg_tok_per_sec"],
        },
    ]

    if "--merge-ext" in sys.argv:
        ext_path = Path(sys.argv[sys.argv.index("--merge-ext") + 1])
        leaderboard.extend(json.loads(ext_path.read_text()))
    else:
        for spec in EXTERNAL_MODELS:
            leaderboard.append(evaluate_external_hf_model(spec, dev))

    payload = {
        "timestamp_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "elapsed_sec": round(time.time() - t0, 2),
        "leaderboard": leaderboard,
        "phobos_mc_final": mc_final,
        "phobos_mc_base": mc_base,
    }
    MULTI_OUT_JSON.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
    for row in leaderboard:
        print(f"{row['model_id']:38s} | Active={row['active_params_m']:6.2f}M | MC20={row['mc20_acc_pct']:5.1f}% | Gen25={row['gen25_score_pct']:5.1f}% | Combined={row['combined_avg_pct']:5.1f}%", flush=True)
    print(f"=== SAVED MULTI-MODEL LEADERBOARD TO {MULTI_OUT_JSON} IN {payload['elapsed_sec']}s ===", flush=True)


if __name__ == "__main__":
    main()