"""Deterministic evaluation: RL (greedy mean action) vs classical.

Layouts are FIXED and identical for every policy/condition:
  single-agent: seeds 1000..1019 (10 general + 10 trap, alternating)
  4-agent:      seeds 2000..2011 (6 general + 6 trap, alternating)
Conditions (same noise RNG seed per layout/condition for both policies):
  clean, obs_noise 0.05, obs_noise 0.10, disturbance 1.0 (accel noise),
  obstacle_shift 0.5 (circles +x 0.5, segments +x 0.5, clipped to arena).
Ablation: 4-agent clean, RL with neighbour-observation components
  (obs indices 11,12,13) zeroed at evaluation time.
Metrics per policy x condition x agent-count, aggregated over agents:
  success_rate, collision_rate, collision_obstacle_rate,
  collision_agent_rate, timeout_rate, mean steps-to-goal (successes),
  mean path length, mean min clearance (successes),
  safety_violation_rate (steps with clearance < 0.15 / active steps).
Trap/general breakdowns are recorded separately in results.json.
Outputs: results.json, eval.csv (one row per agent-episode).
"""
import argparse
import copy
import csv
import json

import numpy as np
import torch

from sim import ARENA, AGENT_R, SAFETY_MARGIN, generate_layout, NavSim
from policies import ActorCritic, classical_action, rl_action


def shifted_layout(layout, shift=0.5):
    lay = copy.deepcopy(layout)
    lay["circles"][:, 0] = np.clip(lay["circles"][:, 0] + shift, 0.8, 11.2)
    seg = lay["segments"].copy()
    seg[:, 0] = np.clip(seg[:, 0] + shift, 0.3, 11.7)
    seg[:, 2] = np.clip(seg[:, 2] + shift, 0.3, 11.7)
    lay["segments"] = seg
    return lay


def run_episode(policy, net, layout, n_agents, obs_noise, disturbance,
                noise_seed, zero_neighbour=False, shift=0.0):
    if shift:
        layout = shifted_layout(layout, shift)
    sim = NavSim(n_agents, seed=noise_seed)
    obs = sim.reset(layout, obs_noise=obs_noise, disturbance=disturbance)
    while True:
        if policy == "classical":
            act = classical_action(obs)
            if obs_noise > 0:
                # classical uses the same noisy observation as RL
                pass
        elif policy == "rl":
            o = obs.copy()
            if zero_neighbour:
                o[:, 11:14] = 0.0
            act = rl_action(net, o, deterministic=True)
        elif policy == "rl_no_neighbour":
            o = obs.copy()
            o[:, 11:14] = 0.0
            act = rl_action(net, o, deterministic=True)
        else:
            raise ValueError(policy)
        obs, _, done, _ = sim.step(act)
        if done:
            break
    core = sim.core
    rows = []
    for a in range(n_agents):
        kind = core.collision_kind[0][a]
        rows.append({
            "n_agents": n_agents,
            "policy": policy,
            "trap": bool(layout["trap"]),
            "success": bool(core.success[0, a]),
            "collision": bool(core.collision[0, a]),
            "collision_obstacle": kind in ("obstacle", "obstacle+agent"),
            "collision_agent": kind in ("agent", "obstacle+agent"),
            "timeout": bool(not core.success[0, a] and not core.collision[0, a]),
            "steps": int(core.steps[0]) if core.success[0, a] else None,
            "steps_total": int(core.steps[0]),
            "path_length": float(core.path_len[0, a]),
            "min_clearance": float(core.min_clearance[0, a]),
            "viol_steps": int(core.viol_steps[0, a]),
            "active_steps": int(core.active_steps[0, a]),
        })
    return rows


def aggregate(rows):
    n = len(rows)
    succ = [r for r in rows if r["success"]]
    return {
        "n_agent_episodes": n,
        "success_rate": float(np.mean([r["success"] for r in rows])),
        "collision_rate": float(np.mean([r["collision"] for r in rows])),
        "collision_obstacle_rate": float(np.mean([r["collision_obstacle"] for r in rows])),
        "collision_agent_rate": float(np.mean([r["collision_agent"] for r in rows])),
        "timeout_rate": float(np.mean([r["timeout"] for r in rows])),
        "mean_steps_to_goal_success": float(np.mean([r["steps"] for r in succ])) if succ else None,
        "mean_path_length": float(np.mean([r["path_length"] for r in rows])),
        "mean_min_clearance_success": float(np.mean([r["min_clearance"] for r in succ])) if succ else None,
        "safety_violation_rate": float(sum(r["viol_steps"] for r in rows) / max(sum(r["active_steps"] for r in rows), 1)),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default="ppo_policy.pt")
    ap.add_argument("--out", default="results.json")
    ap.add_argument("--csv", default="eval.csv")
    args = ap.parse_args()
    net = ActorCritic()
    net.load_state_dict(torch.load(args.weights, map_location="cpu", weights_only=True))
    net.eval()

    conditions = [
        ("clean", dict(obs_noise=0.0, disturbance=0.0, shift=0.0)),
        ("obs_noise_0.05", dict(obs_noise=0.05, disturbance=0.0, shift=0.0)),
        ("obs_noise_0.10", dict(obs_noise=0.10, disturbance=0.0, shift=0.0)),
        ("disturbance_1.0", dict(obs_noise=0.0, disturbance=1.0, shift=0.0)),
        ("obstacle_shift_0.5", dict(obs_noise=0.0, disturbance=0.0, shift=0.5)),
    ]
    all_rows = []
    results = {"protocol": {
        "single_agent_layout_seeds": list(range(1000, 1020)),
        "four_agent_layout_seeds": list(range(2000, 2012)),
        "layout_rule": "alternating general/trap starting with general; same layouts for every policy and condition",
        "noise_seed_rule": "90000 + layout_index, same seed per layout/condition for both policies",
        "weights": args.weights,
        "max_steps": 200, "dt": 0.1, "arena": 12.0,
        "safety_margin": SAFETY_MARGIN,
    }, "conditions": {}}

    for cond_name, cond in conditions:
        results["conditions"][cond_name] = {}
        for n_agents, seeds in ((1, range(1000, 1020)), (4, range(2000, 2012))):
            layouts = []
            for i, s in enumerate(seeds):
                rng = np.random.default_rng(s)
                layouts.append(generate_layout(rng, n_agents, trap=(i % 2 == 1)))
            key = f"{n_agents}_agent"
            results["conditions"][cond_name][key] = {}
            for policy in ("classical", "rl"):
                rows = []
                for li, lay in enumerate(layouts):
                    rows.extend(run_episode(policy, net, lay, n_agents,
                                            cond["obs_noise"], cond["disturbance"],
                                            noise_seed=90000 + li, shift=cond["shift"]))
                for r in rows:
                    r["condition"] = cond_name
                all_rows.extend(rows)
                results["conditions"][cond_name][key][policy] = aggregate(rows)
                for trap_val, label in ((False, "general"), (True, "trap")):
                    sub = [r for r in rows if r["trap"] == trap_val]
                    results["conditions"][cond_name][key][policy][f"breakdown_{label}"] = aggregate(sub)

    # ablation: 4-agent clean, RL with neighbour observation zeroed
    abl_rows, abl_base = [], []
    for i, s in enumerate(range(2000, 2012)):
        rng = np.random.default_rng(s)
        lay = generate_layout(rng, 4, trap=(i % 2 == 1))
        abl_rows.extend(run_episode("rl_no_neighbour", net, lay, 4, 0.0, 0.0, 90000 + i))
        abl_base.extend(run_episode("rl", net, lay, 4, 0.0, 0.0, 90000 + i))
    for r in abl_rows:
        r["condition"] = "ablation_no_neighbour_4agent_clean"
    all_rows.extend(abl_rows)
    results["ablation_no_neighbour_4agent_clean"] = {
        "rl_with_neighbour": aggregate(abl_base),
        "rl_no_neighbour": aggregate(abl_rows),
    }

    with open(args.out, "w") as f:
        json.dump(results, f, indent=2)
    with open(args.csv, "w", newline="") as f:
        fields = ["condition", "policy", "n_agents", "trap", "success", "collision",
                  "collision_obstacle", "collision_agent", "timeout", "steps",
                  "steps_total", "path_length", "min_clearance", "viol_steps", "active_steps"]
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for r in all_rows:
            w.writerow({k: r[k] for k in fields})
    print(json.dumps(results["conditions"], indent=2))
    print("ablation:", json.dumps(results["ablation_no_neighbour_4agent_clean"], indent=2))


if __name__ == "__main__":
    main()