#!/usr/bin/env python3
"""Animated overview video (mp4): AIP-AD swarm decision process + RL training.

Scenes: (1) title, (2) per-view scores of a sample & poolers competing,
(3) GRPO policy evolution over epochs, (4) continual learning curve,
(5) summary board. Rendered with matplotlib FuncAnimation -> ffmpeg.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation, FFMpegWriter
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)
RES = os.path.join(ROOT, "results")
FIG = os.path.join(RES, "figures")
OUT = os.path.join(RES, "aipad_overview.mp4")

with open(os.path.join(RES, "results.json")) as f:
    R = json.load(f)
table = R["fixed_strategies"]
arb = R["grpo_arbiter"]
cats = [c for c in table["max"] if c != "macro"]
snames = ["max", "mean", "median", "top2", "trimmed", "aip_d05", "aip_c03", "aip_soft2", "aip_mt01", "aip_mt025", "aip_mt05", "aip_mt075", "aip_wt05", "aip_wt075"]

# pick the first anomalous sample of the first category for the demo panel
demo_cat = cats[0]
ev = np.load(os.path.join(RES, "scores", f"{demo_cat}.npz.eval.npz"))
vs_all = ev["view_scores"].astype(np.float64)
y_all = ev["view_labels"].max(1)
demo_i = int(np.where(y_all == 1)[0][0])
demo_scores = vs_all[demo_i]

from src.fusion import STRATEGIES  # noqa  (repo root on sys.path via ROOT)
fused_demo = {n: fn(demo_scores) for n, fn in STRATEGIES.items()}

T = 40  # frames per scene


def setup_axis(ax):
    ax.set_axis_off()


fig = plt.figure(figsize=(12.8, 7.2), dpi=100)
fig.patch.set_facecolor("#0b1220")

# ---------- scene frames ----------
ax_title = fig.add_axes([0, 0.86, 1, 0.14])
ax_sc = fig.add_axes([0.06, 0.42, 0.38, 0.38])
ax_pool = fig.add_axes([0.52, 0.42, 0.42, 0.38])
ax_policy = fig.add_axes([0.52, 0.08, 0.42, 0.28])
ax_curve = fig.add_axes([0.06, 0.08, 0.38, 0.28])
for a in (ax_title, ax_sc, ax_pool, ax_policy, ax_curve):
    a.set_facecolor("#0b1220")
    for s in a.spines.values():
        s.set_color("#334155")
    a.tick_params(colors="#cbd5e1", labelsize=8)
    [t.set_color("#cbd5e1") for t in a.get_xticklabels() + a.get_yticklabels()]

for a in (ax_title,):
    a.axis("off")


def draw_static():
    ax_title.clear(); ax_title.axis("off")
    ax_title.text(0.5, 0.62, "AIP-AD · Liars Are Information for Multi-View Industrial Anomaly Detection",
                  ha="center", va="center", color="#f8fafc", fontsize=15, weight="bold")
    ax_title.text(0.5, 0.22, "5 camera-agents report per-view anomalies  →  Byzantine-robust + AIP pooling  →  GRPO arbiter decides",
                  ha="center", va="center", color="#94a3b8", fontsize=10)


draw_static()
ax_sc.set_title(f"Per-view agent reports — '{demo_cat}' sample (anomalous)",
                color="#e2e8f0", fontsize=10, weight="bold")
ax_pool.set_title("Fusion strategies compete for the decision", color="#e2e8f0", fontsize=10, weight="bold")
ax_policy.set_title("GRPO policy over strategies (training epochs)", color="#e2e8f0", fontsize=10, weight="bold")
ax_curve.set_title("Continual learning: macro S-AUROC on all categories", color="#e2e8f0", fontsize=10, weight="bold")


def scene_bars(fr):
    ax_sc.clear()
    ax_sc.set_facecolor("#0b1220")
    for s in ax_sc.spines.values():
        s.set_color("#334155")
    ax_sc.tick_params(colors="#cbd5e1", labelsize=8)
    views = np.arange(5)
    shown = np.clip(demo_scores, None, None)
    cols = ["#ef4444" if v == shown.max() else "#38bdf8" for v in shown]
    ax_sc.bar(views, shown * (fr / T), color=cols, width=0.62)
    med = np.median(shown) * (fr / T)
    ax_sc.axhline(med, color="#facc15", ls="--", lw=1.2)
    ax_sc.text(4.45, med, "swarm consensus (median)", color="#facc15", fontsize=7, va="bottom", ha="right")
    ax_sc.set_xticks(views)
    ax_sc.set_xticklabels([f"cam {i+1}" for i in range(5)])
    ax_sc.set_ylim(0, max(shown.max() * 1.25, 0.1))
    ax_sc.set_ylabel("anomaly score", color="#cbd5e1")
    ax_sc.set_title(f"Per-view agent reports — '{demo_cat}' sample (anomalous)",
                    color="#e2e8f0", fontsize=10, weight="bold")


def scene_pool(fr):
    ax_pool.clear()
    ax_pool.set_facecolor("#0b1220")
    for s in ax_pool.spines.values():
        s.set_color("#334155")
    ax_pool.tick_params(colors="#cbd5e1", labelsize=7)
    keys = ["max", "mean", "median", "top2", "aip_mt01", "aip_wt05", "aip_mt025"]
    vals = np.array([fused_demo[k] for k in keys])
    grow = np.clip(fr / (T * 0.5), 0.02, 1)
    cols = ["#94a3b8", "#94a3b8", "#94a3b8", "#94a3b8", "#f59e0b", "#f59e0b", "#f59e0b"]
    ax_pool.barh(np.arange(len(keys))[::-1], vals * grow, color=cols, height=0.6)
    ax_pool.set_yticks(np.arange(len(keys))[::-1])
    ax_pool.set_yticklabels(keys)
    ax_pool.set_xlim(0, max(vals.max() * 1.3, 0.1))
    ax_pool.set_xlabel("fused anomaly score", color="#cbd5e1")
    ax_pool.set_title("Fusion strategies compete for the decision", color="#e2e8f0", fontsize=10, weight="bold")


def scene_policy(fr):
    ax_policy.clear()
    ax_policy.set_facecolor("#0b1220")
    for s in ax_policy.spines.values():
        s.set_color("#334155")
    ax_policy.tick_params(colors="#cbd5e1", labelsize=7)
    hist = arb["train_history"]
    P = np.array(hist[-1]["mean_probs"])
    frac = fr / T
    idx = min(int(frac * (len(hist) - 1)), len(hist) - 1)
    Pt = np.array(hist[idx]["mean_probs"])
    order = np.argsort(P)
    ax_policy.barh(np.array(snames)[order], Pt[order], color="#dc2626", alpha=0.85, height=0.6)
    ax_policy.set_yticks(np.arange(len(snames))[::-1] * 0 + np.arange(len(snames)))
    ax_policy.set_yticklabels([snames[i] for i in order], fontsize=7)
    ax_policy.set_xlim(0, max(P.max() * 1.25, 0.05))
    ent = hist[idx]["policy_entropy"]
    ax_policy.set_xlabel(f"epoch {hist[idx]['epoch']} · policy entropy {ent:.2f} nats", color="#cbd5e1")
    ax_policy.set_title("GRPO policy over strategies (training epochs)", color="#e2e8f0", fontsize=10, weight="bold")


def scene_curve(fr):
    ax_curve.clear()
    ax_curve.set_facecolor("#0b1220")
    for s in ax_curve.spines.values():
        s.set_color("#334155")
    ax_curve.tick_params(colors="#cbd5e1", labelsize=8)
    run = R["continual"][0]["curve"]
    xs = np.arange(1, len(run) + 1)
    ys = np.array([c["macro_sauroc_all6"] for c in run])
    mmax = np.mean([table["max"][c] for c in cats])
    k = max(1, int((fr / T) * len(xs)))
    ax_curve.plot(xs[:k], ys[:k], "o-", color="#2563eb", lw=2)
    ax_curve.axhline(mmax, color="#94a3b8", ls=":", lw=1.4)
    ax_curve.text(len(xs) * 0.55, mmax, "static max-pool (paper protocol)", color="#94a3b8", fontsize=8, va="bottom")
    ax_curve.set_xticks(xs)
    ax_curve.set_xticklabels([c["after"] for c in run], rotation=25, ha="right", fontsize=7)
    ax_curve.set_ylim(0.5, 1.0)
    ax_curve.set_ylabel("macro S-AUROC", color="#cbd5e1")
    ax_curve.set_title("Continual learning: macro S-AUROC on all categories", color="#e2e8f0", fontsize=10, weight="bold")


SCENES = [(scene_bars, scene_pool), (scene_policy, scene_curve)]
total_frames = T * 2


def update(fr):
    i = fr // T
    fr2 = fr % T
    scene_bars(fr2) if i == 0 else scene_policy(fr2)
    scene_pool(fr2) if i == 0 else scene_curve(fr2)
    draw_static()


anim = FuncAnimation(fig, update, frames=total_frames, interval=110, blit=False)
try:
    writer = FFMpegWriter(fps=9, bitrate=2200)
    anim.save(OUT, writer=writer)
    print("video saved:", OUT, os.path.getsize(OUT))
except Exception as e:
    print("ffmpeg failed:", e)
    OUT_GIF = OUT.replace(".mp4", ".gif")
    anim.save(OUT_GIF, writer="pillow", fps=8)
    print("gif saved:", OUT_GIF, os.path.getsize(OUT_GIF))