#!/usr/bin/env python
"""
Analyze Phase A results and compare with baseline.
Generate comprehensive report for A* paper.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

def load_eval_json(path: Path) -> dict:
    """Load evaluation JSON with error handling."""
    if not path.exists():
        return
    with open(path) as f:
        return json.load(f)

def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Analyze Phase A results for A* paper"
    )
    parser.add_argument("--baseline", type=Path, required=True,
                       help="Baseline experiment directory")
    parser.add_argument("--large-model", type=Path, required=True,
                       help="Phase A2 large model directory")
    parser.add_argument("--hparam-sweep", type=Path, required=False,
                       help="Phase A4 hyperparameter sweep directory")
    parser.add_argument("--horizon-sweep", type=Path, required=False,
                       help="Phase A5 horizon sweep directory")
    parser.add_argument("--out", type=Path, required=True,
                       help="Output JSON path")

    args = parser.parse_args(argv)

    print("=" * 70)
    print("Phase A Results Analysis")
    print("=" * 70)
    print()

    # Load baseline results (3 seeds)
    baseline_successes = []
    for seed in range(3):
        eval_path = args.baseline / f"seed_{seed}" / "policy_rollout.json"
        data = load_eval_json(eval_path)
        if "policy_rollout_success_rate" in data:
            baseline_successes.append(data["policy_rollout_success_rate"])

    baseline_mean = np.mean(baseline_successes) if baseline_successes else 0.2967
    baseline_std = np.std(baseline_successes) if len(baseline_successes) > 1 else 0.0018

    print(f"📊 Baseline Results (current)")
    print(f"   Policy success: {baseline_mean:.4f} ± {baseline_std:.4f}")
    print(f"   Seeds: {len(baseline_successes)}")
    print()

    # Load Phase A2 large model results
    large_successes = []
    for seed in range(3):
        eval_path = args.large_model / f"seed_{seed}" / "policy_rollout.json"
        data = load_eval_json(eval_path)
        if "policy_rollout_success_rate" in data:
            large_successes.append(data["policy_rollout_success_rate"])

    if large_successes:
        large_mean = np.mean(large_successes)
        large_std = np.std(large_successes) if len(large_successes) > 1 else 0.0
        improvement = large_mean - baseline_mean
        relative_improvement = (improvement / baseline_mean) * 100

        print(f"🚀 Phase A2: Large Model Results")
        print(f"   Policy success: {large_mean:.4f} ± {large_std:.4f}")
        print(f"   Improvement: {improvement:+.4f} ({relative_improvement:+.1f}%)")
        print(f"   Seeds: {len(large_successes)}")

        if large_mean >= 0.40:
            print(f"   ✅ Target 40%+ ACHIEVED!")
        else:
            print(f"   ⚠️  Target 40%+ not yet reached (need +{0.40 - large_mean:.4f})")
        print()
    else:
        print("⚠️  No Phase A2 results found yet")
        print()
        large_mean = baseline_mean
        large_std = baseline_std

    # Compile results
    results = {
        "baseline": {
            "policy_success_mean": float(baseline_mean),
            "policy_success_std": float(baseline_std),
            "seeds": baseline_successes
        },
        "phase_a2_large_model": {
            "policy_success_mean": float(large_mean),
            "policy_success_std": float(large_std),
            "seeds": large_successes,
            "improvement_absolute": float(large_mean - baseline_mean),
            "improvement_relative_pct": float(((large_mean - baseline_mean) / baseline_mean) * 100),
            "target_40pct_achieved": large_mean >= 0.40
        },
        "best_policy_success": float(large_mean),
        "target_success": 0.40,
        "status": "achieved" if large_mean >= 0.40 else "in_progress"
    }

    # 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"✅ Analysis saved to: {args.out}")
    print()

    # Summary
    print("=" * 70)
    print("PHASE A SUMMARY")
    print("=" * 70)
    print()
    print(f"Best result: {large_mean:.1%} policy success")
    print(f"Target: 40%")
    print(f"Status: {'✅ ACHIEVED' if large_mean >= 0.40 else '⏳ IN PROGRESS'}")
    print()

    if large_mean < 0.40:
        gap = 0.40 - large_mean
        print(f"To reach target:")
        print(f"  Need: +{gap:.1%} absolute improvement")
        print(f"  Options:")
        print(f"    1. Longer training (more epochs)")
        print(f"    2. Better hyperparameters (check A4 results)")
        print(f"    3. Longer action horizons (check A5 results)")
        print(f"    4. More data (generate 15-20K groups)")
        print()

    return 0

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