"""
Single-episode visualization for the herding policy.

Outputs (all saved to --out-dir):
  trajectory.png  — full field view: dog path + every sheep path
  timeseries.png  — radius, per-sheep pen distance, action magnitude, reward
  episode.gif     — animated replay (slow enough to read)

Run with no model to watch a RANDOM policy (useful baseline):
  python visualize.py --random --n-sheep 3 --out-dir vis_random/

Usage:
  python visualize.py \\
      --model runs/ppo_consolidation/final_model.zip \\
      --vecnorm runs/ppo_consolidation/vecnorm.pkl \\
      --n-sheep 3 --out-dir vis_out/
"""

import argparse
import os
import math
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.animation as animation
from matplotlib.collections import LineCollection
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
from herding_env import HerdingEnv


# ── colours ──────────────────────────────────────────────────────────────────
SHEEP_COLORS = [
    "#e41a1c", "#377eb8", "#4daf4a", "#984ea3", "#ff7f00",
    "#a65628", "#f781bf", "#999999", "#66c2a5", "#fc8d62",
]
DOG_COLOR   = "#4e342e"
PEN_COLOR   = "#ffe082"
FIELD_COLOR = "#dcedc8"


def make_env(n_sheep, max_steps, seed=42):
    def _init():
        env = HerdingEnv(n_sheep=n_sheep, max_steps=max_steps)
        env.reset(seed=seed)
        return env
    return _init


def run_episode(model, env, n_sheep, max_steps):
    """Run one deterministic episode; return recorded history."""
    obs      = env.reset()
    inner    = env.envs[0]
    done     = False

    dog_xs, dog_ys   = [], []
    sheep_xs         = [[] for _ in range(n_sheep)]
    sheep_ys         = [[] for _ in range(n_sheep)]
    radii            = []
    pen_dists        = [[] for _ in range(n_sheep)]
    action_mags      = []
    rewards          = []
    penned_at        = [None] * n_sheep   # step when each sheep was penned

    step = 0
    while not done:
        if model is None:
            action = env.action_space.sample()[np.newaxis]
        else:
            action, _ = model.predict(obs, deterministic=True)

        obs, reward, dones, infos = env.step(action)
        done = dones[0]
        step += 1

        dx, dy = float(inner.dog_pos[0]), float(inner.dog_pos[1])
        dog_xs.append(dx); dog_ys.append(dy)

        com, radius, _ = inner._flock_stats()
        radii.append(radius)
        rewards.append(float(reward[0]))

        act = action[0]
        action_mags.append(float(np.linalg.norm(act)))

        for i in range(n_sheep):
            sx, sy = float(inner.sheep_pos[i][0]), float(inner.sheep_pos[i][1])
            sheep_xs[i].append(sx)
            sheep_ys[i].append(sy)
            pen_dists[i].append(float(np.linalg.norm(inner.sheep_pos[i] - inner.PEN_CENTER)))
            if inner.penned[i] and penned_at[i] is None:
                penned_at[i] = step

    info = infos[0]
    n_penned = info.get("n_penned", 0)
    success  = n_penned == n_sheep

    return dict(
        dog_xs=dog_xs, dog_ys=dog_ys,
        sheep_xs=sheep_xs, sheep_ys=sheep_ys,
        radii=radii, pen_dists=pen_dists,
        action_mags=action_mags, rewards=rewards,
        penned_at=penned_at,
        n_penned=n_penned, n_sheep=n_sheep,
        success=success, steps=step,
    )


# ── plot helpers ─────────────────────────────────────────────────────────────

def draw_field(ax):
    ax.set_xlim(-16, 16); ax.set_ylim(-16, 16)
    ax.set_aspect("equal"); ax.set_facecolor(FIELD_COLOR)
    ax.add_patch(mpatches.Rectangle((-15,-15), 30, 30,
                 fill=False, edgecolor="#795548", lw=2))
    ax.add_patch(mpatches.Rectangle((10,-15), 3, 7,
                 facecolor=PEN_COLOR, edgecolor="#795548", lw=2))
    ax.text(11.5, -11.5, "pen", ha="center", va="center",
            fontsize=8, color="#795548")


def faded_path(ax, xs, ys, color, lw=1.5, label=None):
    """Draw a path with alpha fading from start (transparent) to end (opaque)."""
    n = len(xs)
    if n < 2:
        return
    points  = np.array([xs, ys]).T.reshape(-1, 1, 2)
    segs    = np.concatenate([points[:-1], points[1:]], axis=1)
    alphas  = np.linspace(0.15, 1.0, len(segs))
    colors  = [(*matplotlib.colors.to_rgb(color), a) for a in alphas]
    lc = LineCollection(segs, colors=colors, linewidth=lw)
    ax.add_collection(lc)
    if label:
        ax.plot([], [], color=color, lw=lw, label=label)


# ── main plots ────────────────────────────────────────────────────────────────

def plot_trajectory(hist, out_path):
    fig, ax = plt.subplots(figsize=(7, 7))
    draw_field(ax)

    # Sheep paths
    for i in range(hist["n_sheep"]):
        c = SHEEP_COLORS[i % len(SHEEP_COLORS)]
        xs, ys = hist["sheep_xs"][i], hist["sheep_ys"][i]
        faded_path(ax, xs, ys, c, lw=1.2, label=f"sheep {i+1}")
        ax.plot(xs[0], ys[0], "o", color=c, ms=7, zorder=4)
        pa = hist["penned_at"][i]
        end = pa if pa is not None else -1
        ax.plot(xs[end], ys[end], "*", color=c, ms=11, zorder=5)

    # Dog path
    faded_path(ax, hist["dog_xs"], hist["dog_ys"], DOG_COLOR, lw=2.0, label="dog")
    ax.plot(hist["dog_xs"][0], hist["dog_ys"][0], "s", color=DOG_COLOR, ms=10, zorder=5)
    ax.plot(hist["dog_xs"][-1], hist["dog_ys"][-1], "D", color=DOG_COLOR, ms=10, zorder=5)

    result = "SUCCESS" if hist["success"] else f"FAIL ({hist['n_penned']}/{hist['n_sheep']} penned)"
    ax.set_title(f"Trajectory — {result} — {hist['steps']} steps", fontsize=12)
    ax.legend(loc="upper left", fontsize=8)
    plt.tight_layout()
    fig.savefig(out_path, dpi=120)
    plt.close(fig)
    print(f"  saved {out_path}")


def plot_timeseries(hist, out_path):
    t      = np.arange(hist["steps"])
    fig, axes = plt.subplots(4, 1, figsize=(12, 10), sharex=True)

    # 1. Flock radius
    axes[0].plot(t, hist["radii"], color="steelblue")
    axes[0].axhline(5.0, color="orange", ls="--", lw=1, label="compact threshold (5m)")
    axes[0].set_ylabel("flock radius (m)")
    axes[0].legend(fontsize=8)
    axes[0].set_title("Flock radius — goal: get below 5m")

    # 2. Per-sheep distance to pen
    for i in range(hist["n_sheep"]):
        c = SHEEP_COLORS[i % len(SHEEP_COLORS)]
        axes[1].plot(t, hist["pen_dists"][i], color=c, lw=1, label=f"sheep {i+1}")
        pa = hist["penned_at"][i]
        if pa is not None:
            axes[1].axvline(pa, color=c, ls=":", lw=1)
    axes[1].set_ylabel("dist to pen (m)")
    axes[1].legend(fontsize=7, ncol=min(hist["n_sheep"], 5))
    axes[1].set_title("Per-sheep distance to pen — goal: all reach 0")

    # 3. Action magnitude (how fast dog is moving)
    axes[2].plot(t, hist["action_mags"], color="tomato", lw=1)
    axes[2].axhline(1.0, color="gray", ls="--", lw=1, label="max")
    axes[2].set_ylabel("action ||(vx,vy)||")
    axes[2].set_ylim(0, 1.5)
    axes[2].set_title("Dog action magnitude — 0=stopped, 1=full speed")
    axes[2].legend(fontsize=8)

    # 4. Reward per step
    axes[3].plot(t, hist["rewards"], color="purple", lw=1, alpha=0.7)
    axes[3].axhline(0, color="black", lw=0.5)
    axes[3].set_ylabel("reward")
    axes[3].set_xlabel("step")
    axes[3].set_title("Reward per step")

    result = "SUCCESS" if hist["success"] else f"FAIL ({hist['n_penned']}/{hist['n_sheep']} penned)"
    fig.suptitle(f"n_sheep={hist['n_sheep']}  {result}  {hist['steps']} steps", fontsize=13)
    plt.tight_layout()
    fig.savefig(out_path, dpi=120)
    plt.close(fig)
    print(f"  saved {out_path}")


def save_gif(hist, out_path, fps=15, skip=5):
    """Animated replay, every `skip` steps."""
    n    = hist["n_sheep"]
    idxs = list(range(0, hist["steps"], skip))

    fig, ax = plt.subplots(figsize=(6, 6))

    def _frame(k):
        ax.clear()
        draw_field(ax)
        t = idxs[k]

        for i in range(n):
            c  = SHEEP_COLORS[i % len(SHEEP_COLORS)]
            s0 = max(0, t - 30)
            ax.plot(hist["sheep_xs"][i][s0:t+1],
                    hist["sheep_ys"][i][s0:t+1],
                    color=c, lw=0.8, alpha=0.5)
            color = "#ff69b4" if (hist["penned_at"][i] is not None
                                   and t >= hist["penned_at"][i]) else c
            ax.plot(hist["sheep_xs"][i][t], hist["sheep_ys"][i][t],
                    "o", color=color, ms=10, zorder=4,
                    markeredgecolor="#555", markeredgewidth=1)

        s0 = max(0, t - 30)
        ax.plot(hist["dog_xs"][s0:t+1], hist["dog_ys"][s0:t+1],
                color=DOG_COLOR, lw=1.5, alpha=0.6)
        ax.plot(hist["dog_xs"][t], hist["dog_ys"][t],
                "s", color=DOG_COLOR, ms=13, zorder=5,
                markeredgecolor="black", markeredgewidth=1.5)

        r = hist["radii"][t]
        ax.set_title(f"step {t}/{hist['steps']}  radius={r:.1f}m  "
                     f"penned={hist['n_penned'] if t==hist['steps']-1 else '?'}/{n}",
                     fontsize=10)

    ani = animation.FuncAnimation(fig, _frame, frames=len(idxs), interval=1000//fps)
    ani.save(out_path, writer="pillow", fps=fps)
    plt.close(fig)
    print(f"  saved {out_path}")


# ── entry point ───────────────────────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model",     default=None, help="Model .zip (omit for random policy)")
    p.add_argument("--vecnorm",   default=None)
    p.add_argument("--n-sheep",   type=int, default=3)
    p.add_argument("--max-steps", type=int, default=2000)
    p.add_argument("--seed",      type=int, default=42)
    p.add_argument("--out-dir",   default="vis_out")
    p.add_argument("--random",    action="store_true",
                   help="Use random policy (baseline comparison)")
    p.add_argument("--gif-fps",   type=int, default=15)
    p.add_argument("--gif-skip",  type=int, default=5,
                   help="Render every Nth step in the GIF")
    p.add_argument("--no-gif",    action="store_true")
    return p.parse_args()


def main():
    args = parse_args()
    os.makedirs(args.out_dir, exist_ok=True)

    raw = DummyVecEnv([make_env(args.n_sheep, args.max_steps, args.seed)])

    if args.random or args.model is None:
        print("Using RANDOM policy")
        env   = raw
        model = None
    else:
        if args.vecnorm:
            env = VecNormalize.load(args.vecnorm, raw)
            env.training    = False
            env.norm_reward = False
        else:
            env = raw
        model = PPO.load(args.model, env=env)
        print(f"Loaded model: {args.model}")

    print(f"Running episode  n_sheep={args.n_sheep}  seed={args.seed} ...")
    hist = run_episode(model, env, args.n_sheep, args.max_steps)

    result = "SUCCESS" if hist["success"] else f"FAIL ({hist['n_penned']}/{hist['n_sheep']} penned)"
    print(f"Episode done: {result}  steps={hist['steps']}")
    print(f"  min radius : {min(hist['radii']):.2f} m")
    print(f"  mean reward: {np.mean(hist['rewards']):.4f}")
    print(f"  mean action: {np.mean(hist['action_mags']):.3f}")

    env.close()

    plot_trajectory(hist, os.path.join(args.out_dir, "trajectory.png"))
    plot_timeseries(hist, os.path.join(args.out_dir, "timeseries.png"))
    if not args.no_gif:
        save_gif(hist, os.path.join(args.out_dir, "episode.gif"),
                 fps=args.gif_fps, skip=args.gif_skip)

    print(f"\nAll outputs saved to {args.out_dir}/")


if __name__ == "__main__":
    main()
