# --- twoquarks bootstrap (path-stable imports) ---
import sys
from pathlib import Path

_ROOT = Path(__file__).resolve().parents[1]  # project root (sibling of exp/)
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))
# -------------------------------------------------

import argparse
import csv
import os
from pathlib import Path

import numpy as np

# Robust imports regardless of where python is launched from.
_THIS_DIR = Path(__file__).resolve().parent
_QUARK_DIR = _THIS_DIR.parent  # .../down
_PROJECT_DIR = _QUARK_DIR.parent  # .../Down

if str(_QUARK_DIR) not in sys.path:
    sys.path.insert(0, str(_QUARK_DIR))

from envs.env_paradox import EpistemicValleyEnv
from levo.agent_levo_paradox import HFLevoAgent, LevoParadoxIsomerAgent


def run_experiment(
    out_csv: str | Path,
    episodes_per_phase: int = 400,
    seed: int = 2025,
) -> None:
    """
    Run HF-Levo and Levo Paradox tabular agents and log results to CSV.

    The CSV schema is:
        global_episode,phase,episode,env_seed,agent,episode_reward,failure_mode,rho_state
    """
    out_path = Path(out_csv)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    rng = np.random.default_rng(seed)

    # env for shape information only
    probe_env = EpistemicValleyEnv(phase=1, seed=seed)
    n_states = probe_env.n_states
    n_actions = probe_env.n_actions

    agents = [
        HFLevoAgent(n_states, n_actions, seed=seed),
        LevoParadoxIsomerAgent(n_states, n_actions, seed=seed + 1),
    ]
    agent_names = ["HFLevo", "LevoParadoxIsomer"]

    global_ep = 0

    with out_path.open("w", newline="") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "global_episode",
                "phase",
                "episode",
                "env_seed",
                "agent",
                "episode_reward",
                "failure_mode",
                "rho_state",
            ],
        )
        writer.writeheader()

        for phase in (1, 2, 3):
            for local_ep in range(int(episodes_per_phase)):
                env_seed = int(rng.integers(0, 2**31 - 1))

                for agent, name in zip(agents, agent_names, strict=True):
                    env = EpistemicValleyEnv(phase=phase, seed=env_seed)
                    s_idx, _info = env.reset()
                    a = agent.select_action(s_idx)
                    s_next_idx, reward, done, step_info = env.step(a)

                    failure_mode = step_info.get("failure_mode", "neutral")

                    rho_state = (
                        float(getattr(agent, "rho")[s_idx])  # type: ignore[attr-defined]
                        if hasattr(agent, "rho")
                        else float("nan")
                    )

                    agent.update(
                        s_idx=s_idx,
                        a=a,
                        r=reward,
                        s_next_idx=s_next_idx,
                        done=done,
                        failure_mode=failure_mode,
                    )

                    writer.writerow(
                        {
                            "global_episode": global_ep,
                            "phase": phase,
                            "episode": local_ep,
                            "env_seed": env_seed,
                            "agent": name,
                            "episode_reward": float(reward),
                            "failure_mode": failure_mode,
                            "rho_state": rho_state,
                        }
                    )

                global_ep += 1


def _build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Run DOWN paradox tabular experiment and save CSV.")
    p.add_argument("--episodes", type=int, default=400, help="Episodes per phase (phases: 1,2,3).")
    p.add_argument("--seed", type=int, default=2025, help="RNG seed.")
    p.add_argument(
        "--out",
        type=str,
        default="",
        help="Optional output CSV path. If omitted, uses TWOQUARKS_RESULTS_DIR or down/results.",
    )
    return p


if __name__ == "__main__":
    args = _build_argparser().parse_args()

    # Allow a caller (e.g., run_all.py) to route outputs to a shared folder.
    results_dir = Path(os.environ.get("TWOQUARKS_RESULTS_DIR", (_QUARK_DIR / "results").as_posix()))
    results_dir.mkdir(parents=True, exist_ok=True)

    out_csv = Path(args.out) if args.out.strip() else (results_dir / "down_paradox_tabular_results.csv")

    run_experiment(out_csv.as_posix(), episodes_per_phase=int(args.episodes), seed=int(args.seed))