import os
import sys
from pathlib import Path

THIS_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = THIS_DIR.parent
sys.path.append(str(PROJECT_ROOT))

import pandas as pd
import matplotlib.pyplot as plt

def plot_from_csv(csv_path: Path, out_dir: Path, prefix: str) -> None:
    df = pd.read_csv(csv_path)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Episode return
    if {"episode","reward"}.issubset(df.columns):
        ep_ret = df.groupby("episode")["reward"].sum().reset_index()
        plt.figure()
        plt.plot(ep_ret["episode"], ep_ret["reward"])
        plt.xlabel("episode")
        plt.ylabel("episode_return")
        plt.title(f"{prefix} episode return")
        plt.tight_layout()
        plt.savefig(out_dir / f"{prefix.lower()}_episode_return.png", dpi=160)
        plt.close()

    # Stability / diversity if present
    for col in ["stability","diversity","phase","progress","w_t","H_swarm","S_t","T"]:
        if col in df.columns:
            series = df.groupby("episode")[col].mean().reset_index()
            plt.figure()
            plt.plot(series["episode"], series[col])
            plt.xlabel("episode")
            plt.ylabel(col)
            plt.title(f"{prefix} mean {col}")
            plt.tight_layout()
            plt.savefig(out_dir / f"{prefix.lower()}_{col}.png", dpi=160)
            plt.close()

def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", required=True, help="Path to results CSV (relative to project root is ok).")
    ap.add_argument("--out", default="graphics", help="Output directory for plots (relative ok).")
    ap.add_argument("--prefix", default="Strange", help="Plot filename prefix.")
    args = ap.parse_args()

    csv_path = Path(args.csv)
    if not csv_path.is_absolute():
        csv_path = (PROJECT_ROOT / csv_path).resolve()
    out_dir = Path(args.out)
    if not out_dir.is_absolute():
        out_dir = (PROJECT_ROOT / out_dir).resolve()

    if not csv_path.exists():
        raise FileNotFoundError(f"CSV not found: {csv_path}")

    plot_from_csv(csv_path, out_dir, args.prefix)
    print(f"[ok] plots saved to: {out_dir}")

if __name__ == "__main__":
    main()