"""Plots from results.json + live rollouts of the trained policy.

Outputs (plots/):
  trap_trajectories.png  - one wall-trap layout, RL vs classical
  four_agent.png         - one 4-agent layout, RL vs classical (2 panels)
  clean_comparison.png   - bar chart: clean success / collision, RL vs
                           classical, single-agent and 4-agent
  training_curve.png     - training success/collision over iterations
"""
import argparse
import json

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

from sim import generate_layout, NavSim
from policies import ActorCritic, classical_action, rl_action


def draw(ax, layout, traj, goals, title):
    ax.set_xlim(0, 12); ax.set_ylim(0, 12)
    ax.set_aspect("equal"); ax.set_title(title, fontsize=10)
    ax.set_xlabel("x"); ax.set_ylabel("y")
    for (cx, cy, r) in layout["circles"]:
        ax.add_patch(plt.Circle((cx, cy), r, color="0.7", zorder=1))
        ax.add_patch(plt.Circle((cx, cy), r, fill=False, color="k", lw=0.8, zorder=2))
    for (x1, y1, x2, y2) in layout["segments"]:
        ax.plot([x1, x2], [y1, y2], color="k", lw=4, solid_capstyle="butt", zorder=2)
    colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd"]
    for i, tr in enumerate(traj):
        tr = np.array(tr)
        ax.plot(tr[:, 0], tr[:, 1], color=colors[i % 4], lw=1.5, zorder=3)
        ax.scatter([tr[0, 0]], [tr[0, 1]], marker="o", s=45, color=colors[i % 4], zorder=4)
        ax.scatter([goals[i, 0]], [goals[i, 1]], marker="*", s=160, color=colors[i % 4],
                   edgecolor="k", linewidth=0.5, zorder=4)
    ax.add_patch(plt.Rectangle((0, 0), 12, 12, fill=False, color="k", lw=1.2))


def rollout(policy, net, layout, n_agents):
    sim = NavSim(n_agents, seed=0)
    obs = sim.reset(layout)
    while True:
        act = classical_action(obs) if policy == "classical" else rl_action(net, obs, deterministic=True)
        obs, _, done, _ = sim.step(act)
        if done:
            break
    return sim.traj, sim.success, sim.collision


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default="ppo_policy.pt")
    ap.add_argument("--results", default="results.json")
    ap.add_argument("--curve", default="training_curve.csv")
    ap.add_argument("--outdir", default="plots")
    args = ap.parse_args()
    import os
    os.makedirs(args.outdir, exist_ok=True)
    net = ActorCritic()
    net.load_state_dict(torch.load(args.weights, map_location="cpu", weights_only=True))
    net.eval()

    # trap layout: seed 1001 (second single-agent layout, trap)
    rng = np.random.default_rng(1001)
    trap_layout = generate_layout(rng, 1, trap=True)
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
    for ax, policy in zip(axes, ["classical", "rl"]):
        traj, succ, coll = rollout(policy, net, trap_layout, 1)
        draw(ax, trap_layout, traj, trap_layout["goals"],
             f"{policy}: {'SUCCESS' if succ[0] else 'COLLISION' if coll[0] else 'TIMEOUT'}")
    fig.tight_layout()
    fig.savefig(f"{args.outdir}/trap_trajectories.png", dpi=150)
    plt.close(fig)

    # 4-agent layout: seed 2001 (second 4-agent layout, trap)
    rng = np.random.default_rng(2001)
    lay4 = generate_layout(rng, 4, trap=True)
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
    for ax, policy in zip(axes, ["classical", "rl"]):
        traj, succ, coll = rollout(policy, net, lay4, 4)
        draw(ax, lay4, traj, lay4["goals"],
             f"{policy}: {int(succ.sum())}/4 success, {int(coll.sum())} collisions")
    fig.tight_layout()
    fig.savefig(f"{args.outdir}/four_agent.png", dpi=150)
    plt.close(fig)

    # bar chart from results.json
    with open(args.results) as f:
        res = json.load(f)
    clean = res["conditions"]["clean"]
    labels = ["1-agent\nclassical", "1-agent\nRL", "4-agent\nclassical", "4-agent\nRL"]
    succ = [clean["1_agent"]["classical"]["success_rate"],
            clean["1_agent"]["rl"]["success_rate"],
            clean["4_agent"]["classical"]["success_rate"],
            clean["4_agent"]["rl"]["success_rate"]]
    coll = [clean["1_agent"]["classical"]["collision_rate"],
            clean["1_agent"]["rl"]["collision_rate"],
            clean["4_agent"]["classical"]["collision_rate"],
            clean["4_agent"]["rl"]["collision_rate"]]
    x = np.arange(len(labels)); w = 0.35
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.bar(x - w / 2, succ, w, label="success rate")
    ax.bar(x + w / 2, coll, w, label="collision rate")
    ax.set_xticks(x); ax.set_xticklabels(labels)
    ax.set_ylim(0, 1.05); ax.set_ylabel("rate"); ax.set_title("Clean evaluation: RL vs classical")
    ax.legend()
    for i, v in enumerate(succ):
        ax.text(i - w / 2, v + 0.02, f"{v:.2f}", ha="center", fontsize=9)
    for i, v in enumerate(coll):
        ax.text(i + w / 2, v + 0.02, f"{v:.2f}", ha="center", fontsize=9)
    fig.tight_layout()
    fig.savefig(f"{args.outdir}/clean_comparison.png", dpi=150)
    plt.close(fig)

    # training curve (concatenated curves passed in --curve if present)
    try:
        import csv
        iters, s, c = [], [], []
        with open(args.curve) as f:
            for row in csv.DictReader(f):
                if row["success_rate_finished"] and row["success_rate_finished"] != "nan":
                    iters.append(int(row["iteration"])); s.append(float(row["success_rate_finished"]))
                    c.append(float(row["collision_rate_finished"]))
        fig, ax = plt.subplots(figsize=(8, 4))
        ax.plot(iters, s, label="training success rate (finished episodes)")
        ax.plot(iters, c, label="training collision rate")
        ax.set_xlabel("iteration"); ax.set_ylabel("rate"); ax.legend()
        ax.set_title("PPO training curve (last training run only)")
        fig.tight_layout()
        fig.savefig(f"{args.outdir}/training_curve.png", dpi=150)
        plt.close(fig)
    except FileNotFoundError:
        pass
    print("plots written to", args.outdir)


if __name__ == "__main__":
    main()