#!/usr/bin/env python3
"""Launch CFD closed-loop validation for all completed PySR formulas.

Usage:
    conda run -n pycuda_3_10 python src/SR_analysis/validate/launch_pysr_validation.py

This script scans validate/results/pysr_*.json, picks the best formula
from each pair (front+top), and runs CFD closed-loop validation.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
from typing import Dict, List, Optional, Tuple

RESULTS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "results"))

# Each validation job config
# (scene, front_json, top_json, device, steps, validator_script)
VALIDATION_JOBS = [
    # Illusion 1L - best medium complexity
    {
        "scene": "illusion_1L",
        "front_json": os.path.join(RESULTS_DIR, "pysr_illusion_1L_front.json"),
        "top_json": os.path.join(RESULTS_DIR, "pysr_illusion_1L_top.json"),
        "device": 2,
        "steps": 320,
        "validator": "src/SR_analysis/validate/run_closed_loop_illusion.py",
    },
    # Illusion 0.75L
    {
        "scene": "illusion_0.75L",
        "front_json": os.path.join(RESULTS_DIR, "pysr_illusion_0.75L_front.json"),
        "top_json": os.path.join(RESULTS_DIR, "pysr_illusion_0.75L_top.json"),
        "device": 2,
        "steps": 320,
        "validator": "src/SR_analysis/validate/run_closed_loop_illusion.py",
    },
    # Karman mu (12-dim) — best front from phase+mu result
    {
        "scene": "karman_re100",
        "front_json": os.path.join(RESULTS_DIR, "pysr_karman_re100_mu_front.json"),
        "top_json": os.path.join(RESULTS_DIR, "pysr_karman_re100_mu_top.json"),
        "device": 2,
        "steps": 200,
        "validator": "src/SR_analysis/validate/run_closed_loop.py",
    },
]


def verify_formula_files(jobs: List[Dict]) -> List[Dict]:
    """Verify all formula files exist and have valid content."""
    valid = []
    for job in jobs:
        okay = True
        for key in ["front_json", "top_json"]:
            fp = job[key]
            if not os.path.isfile(fp):
                print(f"  MISSING: {fp}")
                okay = False
                continue
            with open(fp) as f:
                data = json.load(f)
            print(f"  {os.path.basename(fp)}: scene={data['scene']}, "
                  f"channel={data['channel']}, best_score={data['best_score']:.4f}")
        if okay:
            valid.append(job)
    return valid


def main():
    print("=" * 60)
    print("PySR Formula Validation Launcher")
    print("=" * 60)
    print("\nChecking formula files:")
    jobs = verify_formula_files(VALIDATION_JOBS)

    if len(jobs) == 0:
        print("\nNo valid formula pairs found. Cannot proceed.")
        sys.exit(1)

    print(f"\nWill launch {len(jobs)} validation jobs (serial, GPU device 2).")
    print("Press Ctrl+C within 5 seconds to abort...")
    import time
    time.sleep(5)

    results = {}
    for job in jobs:
        scene = job["scene"]
        validator = job["validator"]
        front_json = job["front_json"]
        top_json = job["top_json"]
        device = job["device"]
        steps = job["steps"]

        print(f"\n{'=' * 60}")
        print(f"Validating: {scene}")
        print(f"  Front: {os.path.basename(front_json)}")
        print(f"  Top:   {os.path.basename(top_json)}")
        print(f"  GPU:   {device}")
        print(f"{'=' * 60}")

        cmd = [
            "sudo", "-u", "frank14f",
            "conda", "run", "-n", "pycuda_3_10",
            "python", validator,
            "--scene", scene,
            "--device", str(device),
            "--steps", str(steps),
            "--mode", "pysr",
            "--pysr-front", front_json,
            "--pysr-top", top_json,
        ]

        print(f"  Running: {' '.join(cmd)}")
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
            print(result.stdout)
            if result.stderr:
                print(f"  STDERR: {result.stderr[-500:]}")

            # Parse similarity from output
            if result.returncode == 0:
                for line in result.stdout.split("\n"):
                    if "similarity" in line and "action_range" in line:
                        results[scene] = line.strip()
                        break
                    if "similarity (full)" in line:
                        results[scene] = line.strip()
                        break
                if scene not in results:
                    results[scene] = f"return_code={result.returncode}"
            else:
                print(f"  FAILED (code={result.returncode})")
                results[scene] = f"FAILED(code={result.returncode})"
        except subprocess.TimeoutExpired:
            print(f"  TIMEOUT (>1h)")
            results[scene] = "TIMEOUT"
        except Exception as e:
            print(f"  ERROR: {e}")
            results[scene] = f"ERROR({e})"

    # Summary
    print(f"\n{'=' * 60}")
    print("Validation Results Summary")
    print(f"{'=' * 60}")
    for scene, res in results.items():
        print(f"  {scene}: {res}")

    # Save summary
    out_path = os.path.join(RESULTS_DIR, "validation_summary.json")
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved: {out_path}")


if __name__ == "__main__":
    main()
