from __future__ import annotations

import matplotlib
matplotlib.use("Agg")

import argparse
import os
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent

DOWN_DIR = ROOT / "down"
ANTIDOWN_DIR = ROOT / "AntiDown"

# Shared outputs (so plot_dual can find both CSVs)
RESULTS_DIR = ROOT / "results"
GRAPHICS_DIR = ROOT / "graphics"


def _run(cmd: list[str], cwd: Path, env: dict[str, str]) -> None:
    print(f"[run] ({cwd}) $ {' '.join(cmd)}")
    subprocess.run(cmd, cwd=str(cwd), env=env, check=True)


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Run DOWN + AntiDown pipelines and generate plots.")
    p.add_argument("--episodes", type=int, default=200, help="Episodes per phase (for both DOWN and AntiDown).")
    p.add_argument("--seed", type=int, default=2025, help="Seed for DOWN tabular experiment.")
    p.add_argument("--skip_individual_plots", action="store_true", help="Only generate the dual overlay plot.")
    return p


def main() -> None:
    args = _build_parser().parse_args()
    py = sys.executable

    if not DOWN_DIR.exists():
        raise RuntimeError(f"DOWN_DIR not found: {DOWN_DIR} (check folder names)")
    if not ANTIDOWN_DIR.exists():
        raise RuntimeError(f"ANTIDOWN_DIR not found: {ANTIDOWN_DIR} (check folder names)")

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    GRAPHICS_DIR.mkdir(parents=True, exist_ok=True)

    env = os.environ.copy()
    env["TWOQUARKS_RESULTS_DIR"] = str(RESULTS_DIR)
    env["TWOQUARKS_GRAPHICS_DIR"] = str(GRAPHICS_DIR)

    # Canonical CSV names expected by plot_dual_down_antidown.py
    down_csv = RESULTS_DIR / "down_paradox_tabular_results.csv"
    antidown_csv = RESULTS_DIR / "antidown_corrupted_valley_tabular.csv"

    # ----------------------------
    # DOWN (tabular)
    # ----------------------------
    _run(
        [
            py,
            str(DOWN_DIR / "exp" / "run_paradox_tabular.py"),
            "--episodes",
            str(int(args.episodes)),
            "--seed",
            str(int(args.seed)),
            "--out",
            str(down_csv),
        ],
        cwd=DOWN_DIR,
        env=env,
    )

    # ----------------------------
    # AntiDown (tabular)
    # ----------------------------
    # Requires the updated AntiDown runner below (adds --episodes support)
    _run(
        [
            py,
            str(ANTIDOWN_DIR / "exp" / "run_corrupted_valley_tabular.py"),
            "--episodes",
            str(int(args.episodes)),
        ],
        cwd=ANTIDOWN_DIR,
        env=env,
    )

    # ----------------------------
    # Individual plots (optional)
    # ----------------------------
    if not args.skip_individual_plots:
        # If these plot scripts exist in your repo, keep them.
        # They can read from the shared RESULTS_DIR, via explicit --csv.
        plot_down = DOWN_DIR / "exp" / "plot_paradox_results.py"
        if plot_down.exists():
            _run(
                [
                    py,
                    str(plot_down),
                    "--csv",
                    str(down_csv),
                    "--out",
                    str(GRAPHICS_DIR),
                    "--prefix",
                    "Down",
                ],
                cwd=DOWN_DIR,
                env=env,
            )

        plot_antidown = ANTIDOWN_DIR / "exp" / "plot_corrupted_valley_results.py"
        if plot_antidown.exists():
            _run(
                [
                    py,
                    str(plot_antidown),
                    "--csv",
                    str(antidown_csv),
                    "--out",
                    str(GRAPHICS_DIR),
                    "--prefix",
                    "AntiDown",
                ],
                cwd=ANTIDOWN_DIR,
                env=env,
            )

    # ----------------------------
    # Dual overlay plot (canonical)
    # ----------------------------
    _run([py, str(ROOT / "plot_dual_down_antidown.py")], cwd=ROOT, env=env)

    print(f"[ok] Finished.")
    print(f"     CSV: {down_csv}")
    print(f"     CSV: {antidown_csv}")
    print(f"     Plots: {GRAPHICS_DIR}")


if __name__ == "__main__":
    main()