"""Head-to-head benchmark: vanilla Qwen2.5-7B-Instruct vs Prometheus active_model.

ADJUST PATHS

Runs the same 28 probes against both models, captures every prompt/response pair,
and produces a shareable JSON report with no implementation details exposed.

Usage:
    python3 benchmark_comparison.py [--device cuda:0]
"""

import gc
import json
import torch
from datetime import datetime
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoTokenizer
from self_eval import ALL_PROBES, generate, evaluate_probe


VANILLA_DIR = Path("/path/to/models/Qwen2.5-7B-Instruct")
PROMETHEUS_DIR = Path("/path/to/models/tryingforagi/active_model")
OUTPUT_FILE = Path("/path/to/models/prometheus/prometheus_benchmark.json")


def benchmark_model(model_dir, device, label):
    print(f"\n{'='*60}")
    print(f"  Benchmarking: {label}")
    print(f"  Model: {model_dir}")
    print(f"{'='*60}\n")

    model = AutoModelForCausalLM.from_pretrained(
        str(model_dir), device_map={"": device}, torch_dtype="auto",
        low_cpu_mem_usage=True)
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.eval()

    results = []
    by_category = {}

    for i, probe in enumerate(ALL_PROBES):
        response = generate(model, tokenizer, probe["prompt"])
        score, details = evaluate_probe(model, tokenizer, probe)

        entry = {
            "probe": probe["prompt"],
            "category": probe.get("category", "unknown"),
            "response": response,
            "score": round(score, 3),
        }
        results.append(entry)

        cat = entry["category"]
        if cat not in by_category:
            by_category[cat] = {"scores": [], "probes": []}
        by_category[cat]["scores"].append(score)
        by_category[cat]["probes"].append(entry)

        status = "PASS" if score >= 0.5 else "FAIL"
        print(f"  [{i+1:2d}/28] {status} {score:.3f}  {cat:<15s}  {probe['prompt'][:60]}")

    category_scores = {}
    for cat, data in by_category.items():
        avg = sum(data["scores"]) / len(data["scores"])
        category_scores[cat] = round(avg, 3)

    overall = sum(e["score"] for e in results) / len(results)

    del model, tokenizer
    gc.collect()
    torch.cuda.empty_cache()

    return {
        "label": label,
        "overall_score": round(overall, 3),
        "category_scores": dict(sorted(category_scores.items(), key=lambda x: x[1])),
        "probes": results,
    }


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", default="cuda:0")
    args = parser.parse_args()

    report = {
        "title": "Project Prometheus: Before vs After Benchmark",
        "timestamp": datetime.now().isoformat(),
        "description": (
            "Head-to-head comparison of vanilla Qwen2.5-7B-Instruct against "
            "the same model after one autonomous self-improvement cycle. "
            "The model evaluated itself, identified its weakest area, selected "
            "a modification strategy, and validated the result -- all without "
            "human intervention. This benchmark re-runs the same 28 probes "
            "against both models to independently verify the improvement."
        ),
        "base_model": "Qwen/Qwen2.5-7B-Instruct",
        "probes_count": len(ALL_PROBES),
        "categories": sorted(set(p.get("category", "unknown") for p in ALL_PROBES)),
    }

    print("Phase 1: Benchmarking vanilla model...")
    report["vanilla"] = benchmark_model(VANILLA_DIR, args.device, "Vanilla Qwen2.5-7B-Instruct")

    print("\nPhase 2: Benchmarking Prometheus model...")
    report["prometheus"] = benchmark_model(PROMETHEUS_DIR, args.device, "Prometheus (1 self-improvement cycle)")

    v_score = report["vanilla"]["overall_score"]
    p_score = report["prometheus"]["overall_score"]
    delta = round(p_score - v_score, 3)

    print(f"\n{'='*60}")
    print(f"  RESULTS")
    print(f"{'='*60}")
    print(f"  Vanilla:    {v_score}")
    print(f"  Prometheus: {p_score}")
    print(f"  Delta:      {'+' if delta >= 0 else ''}{delta}")
    print(f"\n  Category breakdown:")

    all_cats = sorted(set(
        list(report["vanilla"]["category_scores"].keys()) +
        list(report["prometheus"]["category_scores"].keys())
    ))
    for cat in all_cats:
        vs = report["vanilla"]["category_scores"].get(cat, 0)
        ps = report["prometheus"]["category_scores"].get(cat, 0)
        d = round(ps - vs, 3)
        marker = " <--" if d != 0 else ""
        print(f"    {cat:<15s}  {vs:.3f} -> {ps:.3f}  ({'+' if d >= 0 else ''}{d}){marker}")

    report["summary"] = {
        "vanilla_overall": v_score,
        "prometheus_overall": p_score,
        "delta": delta,
        "improved_categories": [
            cat for cat in all_cats
            if report["prometheus"]["category_scores"].get(cat, 0) >
               report["vanilla"]["category_scores"].get(cat, 0)
        ],
        "degraded_categories": [
            cat for cat in all_cats
            if report["prometheus"]["category_scores"].get(cat, 0) <
               report["vanilla"]["category_scores"].get(cat, 0)
        ],
    }

    with open(OUTPUT_FILE, "w") as f:
        json.dump(report, f, indent=2)
    print(f"\n  Full report saved to {OUTPUT_FILE}")


if __name__ == "__main__":
    main()