#!/usr/bin/env python
"""
Evaluate all Phase A4 hyperparameter configs.
Find best configuration for achieving 40%+ success.
"""
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 A4 hyperparameter configs"
    )
    parser.add_argument("--config-dir", type=Path,
                        default=Path("/scratch/knguy52/dovla/experiments/phase_a4_hparam_sweep"),
                        help="Phase A4 configs 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_a4_results.json"),
                        help="Output JSON path")

    args = parser.parse_args(argv)

    print("=" * 70)
    print("Phase A4 Hyperparameter Evaluation")
    print("=" * 70)
    print(f"Config dir: {args.config_dir}")
    print(f"Dataset: {args.dataset}")
    print()

    configs = sorted([d for d in args.config_dir.iterdir() if d.is_dir()])
    print(f"Found {len(configs)} configs to evaluate")
    print()

    results = {}

    for config_dir in configs:
        config_name = config_dir.name
        checkpoint = config_dir / "best.pt"
        eval_out = config_dir / "lattice_eval.json"

        print(f"Evaluating: {config_name}")

        if not checkpoint.exists():
            print(f"  ⚠️ Checkpoint not found: {checkpoint}")
            results[config_name] = {"status": "missing_checkpoint"}
            continue

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

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

            if eval_out.exists():
                with open(eval_out) as f:
                    data = json.load(f)
                # Extract key metrics
                results[config_name] = {
                    "status": "complete",
                    "selected_success_rate": data.get("selected_success_rate",
                                  data.get("policy_rollout_success_rate", 0)),
                    "top1_action_selection": data.get("top1_action_selection", 0),
                    "pairwise_ranking_accuracy": data.get("pairwise_ranking_accuracy", 0),
                    "selection_regret": data.get("selection_regret", 0),
                    "oracle_success_rate": data.get("oracle_success_rate", 0)
                }
                print(f"    Success: {results[config_name]['selected_success_rate']:.4f}")
            else:
                results[config_name] = {"status": "no_output"}

        except subprocess.CalledProcessError as e:
            print(f"  ✗ Evaluation failed: {e}")
            results[config_name] = {"status": "failed", "error": str(e)}
        except subprocess.TimeoutExpired:
            print(f"  ⏰ Timeout")
            results[config_name] = {"status": "timeout"}

        print()

    # Summary
    print("=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print()

    successful = [r for r in results.values() if r.get("status") == "complete"]
    if successful:
        # Sort by selected success rate
        sorted_results = sorted(successful,
                              key=lambda x: x["selected_success_rate"],
                              reverse=True)

        print("Top configurations:")
        for i, r in enumerate(sorted_results[:5], 1):
            print(f"  {i}. {r['selected_success_rate']:.4f} (top1: {r['top1_action_selection']:.4f})")

        best = sorted_results[0]
        print()
        print(f"Best config: {best['selected_success_rate']:.4f}")
        print(f"Target 40%: {'✅ ACHIEVED' if best['selected_success_rate'] >= 0.40 else '❌ NOT YET'}")

    # Save results
    args.out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w") as f:
        json.dump({
            "configs": results,
            "best_by_success": max(successful, key=lambda x: x["selected_success_rate"]) if successful else None
        }, f, indent=2)

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

    return 0

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