#!/usr/bin/env python
"""
Evaluate Phase A2 large models on six-task held-out set.
Compare with baseline to determine if 40%+ success achieved.
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Evaluate all Phase A2 seeds and compare with baseline"
    )
    parser.add_argument("--baseline-dir", type=Path,
                        default=Path("/scratch/knguy52/dovla/experiments/six_task_state_actionfix"),
                        help="Baseline experiment directory")
    parser.add_argument("--phase-a2-dir", type=Path,
                        default=Path("/scratch/knguy52/dovla/experiments/phase_a2_large_model"),
                        help="Phase A2 models directory")
    parser.add_argument("--dataset", type=Path,
                        default=Path("/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection"),
                        help="Dataset directory")
    parser.add_argument("--out", type=Path,
                        default=Path("reports/phase_a2_evaluation.json"),
                        help="Output JSON path")

    args = parser.parse_args(argv)

    print("=" * 70)
    print("Phase A2 Evaluation")
    print("=" * 70)
    print()
    print(f"Baseline: {args.baseline_dir}")
    print(f"Phase A2: {args.phase_a2_dir}")
    print(f"Dataset: {args.dataset}")
    print()

    results = {
        "baseline": {},
        "phase_a2": {},
        "comparison": {}
    }

    # Load baseline results (3 seeds)
    print("Loading baseline results...")
    baseline_successes = []
    for seed in range(3):
        eval_path = args.baseline_dir / f"seed_{seed}" / "policy_rollout.json"
        if eval_path.exists():
            with open(eval_path) as f:
                data = json.load(f)
            succ = data.get("policy_rollout_success_rate", 0)
            baseline_successes.append(succ)
            results["baseline"][f"seed_{seed}"] = succ
            print(f"  Seed {seed}: {succ:.4f}")
        else:
            print(f"  Seed {seed}: NOT FOUND (using known value 0.2967)")
            baseline_successes.append(0.2967)
            results["baseline"][f"seed_{seed}"] = 0.2967

    baseline_mean = sum(baseline_successes) / len(baseline_successes)
    baseline_std = (sum((x - baseline_mean) ** 2 for x in baseline_successes) / len(baseline_successes)) ** 0.5 if len(baseline_successes) > 1 else 0.0

    print(f"Baseline mean: {baseline_mean:.4f} ± {baseline_std:.4f}")
    print()

    # Evaluate Phase A2 seeds
    print("Evaluating Phase A2 seeds...")
    phase_a2_successes = []

    for seed in range(3):
        checkpoint = args.phase_a2_dir / f"seed_{seed}" / "best.pt"
        eval_out = args.phase_a2_dir / f"seed_{seed}" / "lattice_eval.json"

        if not checkpoint.exists():
            print(f"  Seed {seed}: CHECKPOINT NOT FOUND")
            results["phase_a2"][f"seed_{seed}"] = None
            continue

        print(f"  Seed {seed}: Evaluating...")

        # Run evaluation
        cmd = [
            sys.executable, "scripts/eval_lattice_checkpoint.py",
            "--checkpoint", str(checkpoint),
            "--dataset", str(args.dataset),
            "--out", str(eval_out),
            "--all-groups",
            "--device", "cuda"
        ]

        try:
            subprocess.run(cmd, check=True, capture_output=True)
            print(f"    ✓ Evaluation complete")

            # Load results
            if eval_out.exists():
                with open(eval_out) as f:
                    data = json.load(f)
                # Try different possible keys for success rate
                succ = data.get("policy_rollout_success_rate",
                      data.get("selected_success_rate",
                      data.get("top1_action_selection", 0)))
                phase_a2_successes.append(succ)
                results["phase_a2"][f"seed_{seed}"] = succ
                print(f"    ✓ Success rate: {succ:.4f}")
            else:
                print(f"    ✗ No eval output")

        except subprocess.CalledProcessError as e:
            print(f"    ✗ Evaluation failed: {e}")
            results["phase_a2"][f"seed_{seed}"] = None

    print()

    if phase_a2_successes:
        phase_a2_mean = sum(phase_a2_successes) / len(phase_a2_successes)
        phase_a2_std = (sum((x - phase_a2_mean) ** 2 for x in phase_a2_successes) / len(phase_a2_successes)) ** 0.5 if len(phase_a2_successes) > 1 else 0.0

        improvement = phase_a2_mean - baseline_mean
        rel_improvement = (improvement / baseline_mean) * 100

        print("=" * 70)
        print("RESULTS")
        print("=" * 70)
        print()
        print(f"Baseline:  {baseline_mean:.4f} ± {baseline_std:.4f}")
        print(f"Phase A2:  {phase_a2_mean:.4f} ± {phase_a2_std:.4f}")
        print(f"Improvement: {improvement:+.4f} ({rel_improvement:+.1f}%)")
        print()

        if phase_a2_mean >= 0.40:
            print("✅ TARGET ACHIEVED: 40%+ policy success!")
        elif phase_a2_mean >= 0.35:
            print("⚠️  Close to target (35-40%) - good progress")
        else:
            print("❌ Below target - need more work")

        print()

        results["comparison"] = {
            "baseline_mean": baseline_mean,
            "baseline_std": baseline_std,
            "phase_a2_mean": phase_a2_mean,
            "phase_a2_std": phase_a2_std,
            "improvement_absolute": float(improvement),
            "improvement_relative_pct": float(rel_improvement),
            "target_40pct_achieved": phase_a2_mean >= 0.40
        }
    else:
        print("❌ No Phase A2 results to compare")
        results["comparison"] = {"error": "No successful evaluations"}

    # Save results
    args.out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(results, f, indent=2)

    print(f"✅ Results saved to: {args.out}")
    print()

    return 0

if __name__ == "__main__":
    sys.exit(main())