from __future__ import annotations
import sys
from pathlib import Path

import pandas as pd
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parent
STRANGE_CSV = ROOT / "strange" / "results" / "strange_hypothesis_lab.csv"
ANTISTRANGE_CSV = ROOT / "AntiStrange" / "results" / "antistrange_hypothesis_lab.csv"
OUT_DIR = ROOT / "graphics"

def _episode_return(df: pd.DataFrame) -> pd.DataFrame:
    # Prefer explicit episode_return/return, else sum reward per episode.
    if {"episode", "episode_return"}.issubset(df.columns):
        s = df.groupby("episode")["episode_return"].mean().reset_index()
        s = s.rename(columns={"episode_return": "value"})
        return s
    if {"episode", "return"}.issubset(df.columns):
        s = df.groupby("episode")["return"].mean().reset_index()
        s = s.rename(columns={"return": "value"})
        return s
    if {"episode", "reward"}.issubset(df.columns):
        s = df.groupby("episode")["reward"].sum().reset_index()
        s = s.rename(columns={"reward": "value"})
        return s
    raise KeyError(f"Need columns episode + (episode_return|return|reward). Found: {list(df.columns)}")

def _episode_mean(df: pd.DataFrame, col: str) -> pd.DataFrame | None:
    if {"episode", col}.issubset(df.columns):
        s = df.groupby("episode")[col].mean().reset_index()
        s = s.rename(columns={col: "value"})
        return s
    return None

def main() -> int:
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    if not STRANGE_CSV.exists():
        print(f"[dual-plot] Missing {STRANGE_CSV}")
        return 2
    if not ANTISTRANGE_CSV.exists():
        print(f"[dual-plot] Missing {ANTISTRANGE_CSV}")
        return 2

    df_s = pd.read_csv(STRANGE_CSV)
    df_a = pd.read_csv(ANTISTRANGE_CSV)

    # --- Dual episode return ---
    s_ret = _episode_return(df_s)
    a_ret = _episode_return(df_a)

    plt.figure()
    plt.plot(s_ret["episode"], s_ret["value"], label="Strange")
    plt.plot(a_ret["episode"], a_ret["value"], label="AntiStrange")
    plt.xlabel("episode")
    plt.ylabel("episode_return")
    plt.title("Strange vs AntiStrange — Episode Return")
    plt.legend()
    plt.tight_layout()
    plt.savefig(OUT_DIR / "dual_episode_return.png", dpi=180)
    plt.close()

    # --- Optional dual stability ---
    s_stab = _episode_mean(df_s, "stability")
    a_stab = _episode_mean(df_a, "stability")
    if s_stab is not None or a_stab is not None:
        plt.figure()
        if s_stab is not None:
            plt.plot(s_stab["episode"], s_stab["value"], label="Strange")
        if a_stab is not None:
            plt.plot(a_stab["episode"], a_stab["value"], label="AntiStrange")
        plt.xlabel("episode")
        plt.ylabel("stability (mean per episode)")
        plt.title("Strange vs AntiStrange — Stability")
        plt.legend()
        plt.tight_layout()
        plt.savefig(OUT_DIR / "dual_stability.png", dpi=180)
        plt.close()

    print(f"[dual-plot] Saved plots to {OUT_DIR}")
    return 0

if __name__ == "__main__":
    raise SystemExit(main())