import json, numpy as np, pandas as pd, matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from scipy.stats import t as tdist
import affect_sim as s
plt.rcParams.update({"font.family": "serif", "font.serif": ["Liberation Serif", "DejaVu Serif"],
    "font.size": 7.5, "axes.labelsize": 7.5, "axes.titlesize": 7.5, "legend.fontsize": 6.5,
    "xtick.labelsize": 6.5, "ytick.labelsize": 6.5, "axes.linewidth": 0.6, "lines.linewidth": 0.9,
    "savefig.dpi": 400, "figure.facecolor": "white", "axes.grid": True, "grid.linewidth": 0.3, "grid.alpha": 0.5})
D = "../DATA/"; F = "../FIGURES/"
R = json.load(open(D + "results.json"))
W1 = 3.45
C = {"Face-only KF": "#8c8c8c", "Late fusion": "#c7b37a", "Late fusion+EMA": "#d0743c",
     "Fused KF": "#3b75af", "Gated fused KF": "#1b3f6b"}
def save(fig, n): fig.savefig(F + n + ".png", bbox_inches="tight", pad_inches=0.02); plt.close(fig)

# Fig 1 architecture
fig, ax = plt.subplots(figsize=(W1, 2.35)); ax.set_xlim(0, 10); ax.set_ylim(0, 6.6); ax.axis("off"); ax.grid(False)
def box(x, y, w, h, t, fc="#eef3f9"):
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.02,rounding_size=0.12", fc=fc, ec="#1b3f6b", lw=0.6))
    ax.text(x + w / 2, y + h / 2, t, ha="center", va="center", fontsize=6.2)
def arr(x0, y0, x1, y1): ax.annotate("", (x1, y1), (x0, y0), arrowprops=dict(arrowstyle="-|>", lw=0.6, color="#333", mutation_scale=6))
box(0.1, 5.0, 2.2, 1.0, "Physiological\n(EDA, HR) $y^p$"); box(0.1, 3.6, 2.2, 1.0, "Facial cues\n$y^f$"); box(0.1, 2.2, 2.2, 1.0, "Vocal prosody\n$y^v$")
box(3.0, 2.4, 2.4, 3.4, "Two-timescale\nstate-space\nestimator\n\n$s_t=[x_t;\\,\\mu_t]$\n(emotion; mood)\n\nKalman update +\n$\\chi^2$ gate", fc="#dce7f3")
box(6.0, 4.3, 3.8, 1.4, "Posterior $\\hat{x}_t,\\;P_t$\nconfidence $c_t=e^{-\\mathrm{tr}P^{xx}_t/\\tau}$")
box(6.0, 2.4, 3.8, 1.4, "Constrained behaviour\nmapping $r_t$ (Eq. 9)\nbounds + rate limit")
box(6.0, 0.4, 3.8, 1.3, "Robot display: face,\nprosody, posture", fc="#f4efe4")
for y in (5.5, 4.1, 2.7): arr(2.3, y, 3.0, min(max(y, 2.6), 5.6))
arr(5.4, 5.0, 6.0, 5.0); arr(7.9, 4.3, 7.9, 3.8); arr(7.9, 2.4, 7.9, 1.7)
ax.annotate("", (1.2, 2.2), (6.0, 1.05), arrowprops=dict(arrowstyle="-|>", lw=0.5, ls="--", color="#777", connectionstyle="angle,angleA=180,angleB=90", mutation_scale=6))
ax.text(1.4, 1.25, "human response (closed loop)", fontsize=5.6, color="#555")
save(fig, "fig01_architecture")

# example trial for Fig 2/3/10
rng = np.random.default_rng(314)
x, mu = s.gen_truth(rng); obs = s.gen_obs(rng, x, 1.0, 0.2, 0.05)
o = s.kalman(obs, gate=True, return_P=True); tt = np.arange(s.N) * s.DT
ema = s.late_fusion(obs, ema=0.15)
fig, axs = plt.subplots(3, 1, figsize=(W1, 3.0), sharex=True)
for i, (a, lab) in enumerate(zip(axs, ["Valence", "Arousal", "Dominance"])):
    sd = np.sqrt(o["P"][:, i, i])
    a.fill_between(tt, o["x"][:, i] - 2 * sd, o["x"][:, i] + 2 * sd, color="#3b75af", alpha=0.18, lw=0, label="±2σ")
    a.plot(tt, x[:, i], "k", lw=0.8, label="True (synthetic)")
    a.plot(tt, ema[:, i], color=C["Late fusion+EMA"], lw=0.6, alpha=0.8, label="Late fusion+EMA")
    a.plot(tt, o["x"][:, i], color="#1b3f6b", lw=0.9, label="Gated fused KF")
    a.plot(tt, mu[:, i], ":", color="#555", lw=0.7, label="Mood target $\\mu$")
    av = obs["face"][1] if i == 0 else obs["voice"][1]
    m = ~av; a.fill_between(tt, -1.2, -1.05, where=m, color="#b03a2e", lw=0, step="mid")
    a.set_ylim(-1.2, 1.1); a.set_ylabel(lab)
axs[0].legend(ncol=3, loc="upper center", bbox_to_anchor=(0.5, 1.55), frameon=False)
axs[-1].set_xlabel("Time (s)"); fig.align_ylabels(axs)
save(fig, "fig02_timeseries")

# Fig 3 V-A plane
fig, ax = plt.subplots(figsize=(W1 * 0.62, W1 * 0.62))
th = np.linspace(0, 2 * np.pi, 200); ax.plot(np.cos(th), np.sin(th), color="#bbb", lw=0.5)
ax.axhline(0, color="#999", lw=0.5); ax.axvline(0, color="#999", lw=0.5)
ax.plot(x[::3, 0], x[::3, 1], color="k", lw=0.4, alpha=0.5, label="True")
ax.plot(o["x"][::3, 0], o["x"][::3, 1], color="#1b3f6b", lw=0.5, alpha=0.9, label="Estimate")
for p, n in zip(s.PROTOS, s.PROTO_NAMES):
    if n == "irritated": continue
    n = "anxious / irritated" if n == "anxious" else n
    ax.plot(p[0], p[1], "s", ms=3.5, mfc="#d0743c", mec="k", mew=0.3); ax.text(p[0] + 0.05, p[1] + 0.07, n, fontsize=5.5, bbox=dict(fc="white", ec="none", pad=0.3, alpha=0.8))
ax.set_xlim(-1.1, 1.1); ax.set_ylim(-1.1, 1.1); ax.set_aspect("equal"); ax.set_xlabel("Valence"); ax.set_ylabel("Arousal")
ax.legend(loc="lower right", frameon=False)
save(fig, "fig03_va_plane")

# Fig 4 RMSE bars
fig, ax = plt.subplots(figsize=(W1, 1.8))
dims = ["V", "A", "D", "overall"]; meths = s.METHODS; w = 0.16
for j, m in enumerate(meths):
    vals = [R["main"][m][d][0] for d in dims]; err = [R["main"][m][d][1] for d in dims]
    ax.bar(np.arange(4) + (j - 2) * w, vals, w, yerr=err, color=C[m], label=m, error_kw=dict(lw=0.5, capsize=1.2))
ax.set_xticks(range(4)); ax.set_xticklabels(["Valence", "Arousal", "Dominance", "Overall"]); ax.set_ylabel("RMSE (PAD units)")
ax.set_ylim(0, 0.8); ax.legend(ncol=3, frameon=False, loc="lower center", bbox_to_anchor=(0.5, 1.0))
save(fig, "fig04_rmse")

# Fig 5 dropout sweep
sw = pd.read_csv(D + "dropout_sweep.csv")
fig, ax = plt.subplots(figsize=(W1, 1.7))
for m in ["Face-only KF", "Late fusion+EMA", "Fused KF", "Gated fused KF"]:
    g = sw[sw.method == m].groupby("p").overall; mm = g.mean(); h = tdist.ppf(0.975, 59) * g.std() / np.sqrt(60)
    ax.errorbar(mm.index, mm.values, yerr=h.values, color=C[m], marker="o", ms=2.5, capsize=1.5, lw=0.9, elinewidth=0.5, label=m)
ax.set_xlabel("Stationary dropout probability $p_d$ (per modality)"); ax.set_ylabel("Overall RMSE"); ax.legend(frameon=False, ncol=2)
save(fig, "fig05_dropout")

# Fig 6 NEES
nees = np.load(D + "nees_avg.npy"); nn = R["nees"]
fig, ax = plt.subplots(figsize=(W1, 1.5))
ax.plot(tt, nees, color="#1b3f6b", lw=0.5); ax.axhline(6, color="k", lw=0.6, ls="--", label="Expected (n=6)")
ax.axhspan(nn["lo"], nn["hi"], color="#3b75af", alpha=0.15, label="95% $\\chi^2$ band (50 runs)")
ax.set_xlabel("Time (s)"); ax.set_ylabel("Run-averaged NEES"); ax.set_ylim(3, 10); ax.legend(frameon=False, ncol=2, loc="upper right")
save(fig, "fig06_nees")

# Fig 7 Pareto of standardised effects
eff = R["doe"]["effects"]; lab = {"X1_noise": "A", "X2_dropout": "B", "X3_outlier": "C", "X4_gating": "D"}
names = {k: ":".join(lab[p] for p in k.split(":")).replace(":", "") for k in eff}
items = sorted(eff.items(), key=lambda kv: abs(kv[1]["t"]))
tcrit = tdist.ppf(0.975, R["doe"]["dfres"])
fig, ax = plt.subplots(figsize=(W1, 2.1))
ax.barh([names[k] for k, _ in items], [abs(v["t"]) for _, v in items], color=["#1b3f6b" if abs(v["t"]) > tcrit else "#bbb" for _, v in items], height=0.65)
ax.axvline(tcrit, color="#b03a2e", lw=0.7, ls="--"); ax.text(tcrit + 0.3, 0.2, f"t$_{{0.975}}$={tcrit:.2f}", color="#b03a2e", fontsize=6)
ax.set_xlabel("|Standardised effect| (t)"); ax.text(0.98, 0.04, "A: noise  B: dropout\nC: outliers  D: gating", transform=ax.transAxes, ha="right", fontsize=6)
save(fig, "fig07_pareto")

# Fig 8 interaction plot
dd = pd.read_csv(D + "doe_minitab.csv")
fig, axs = plt.subplots(1, 2, figsize=(W1, 1.55), sharey=True)
for a, ns in zip(axs, (-1, 1)):
    sub = dd[dd.X1_noise == ns]
    for gte, col, lb in ((-1, "#3b75af", "No gating"), (1, "#1b3f6b", "χ² gating")):
        mm = sub[sub.X4_gating == gte].groupby("X3_outlier").Y_rmse.mean()
        a.plot(["0", "0.10"], mm.values, marker="o", ms=3, color=col, ls="-" if gte > 0 else "--", label=lb)
    a.set_title(f"Noise scale = {0.5 if ns < 0 else 1.5}"); a.set_xlabel("Outlier rate $p_o$")
axs[0].set_ylabel("Mean overall RMSE"); axs[0].legend(frameon=False)
save(fig, "fig08_interaction")

# Fig 9 behaviour Pareto
bh = pd.read_csv(D + "behaviour_sweep.csv"); bh = bh[bh.rate == True]
fig, ax = plt.subplots(figsize=(W1, 1.9))
for tau, col, lb in ((None, "#d0743c", "No confidence scaling"), (0.15, "#1b3f6b", "Confidence scaling, τ = 0.15")):
    sub = bh[bh.tau.isna()] if tau is None else bh[bh.tau == tau]
    ax.plot(sub.jitter, sub.track, marker="o", ms=2.8, color=col, label=lb)
    for _, r in sub.iterrows():
        if r.lam in (0, 2, 8, 32): ax.text(r.jitter * 1.08, r.track - (0.006 if tau else -0.002), f"λ={int(r.lam)}", fontsize=5.3, color=col)
b0 = R["behaviour"][0]
ax.axhline(1.35 * b0["track"], color="#777", ls=":", lw=0.7); ax.text(0.03, 1.35 * b0["track"] + 0.002, "1.35× baseline tracking error", fontsize=5.5, color="#555")
sel = bh[(bh.lam == 2) & (bh.tau == 0.15)].iloc[0]
ax.plot(sel.jitter, sel.track, "*", ms=8, mfc="#f2c14e", mec="k", mew=0.4, label="Selected design")
ax.set_xscale("log"); ax.set_xlabel("Display jitter, RMS $\\Delta^2 r/\\Delta t^2$ (s$^{-2}$)"); ax.set_ylabel("Tracking RMSE vs. oracle")
ax.legend(frameon=False, fontsize=6, loc="upper right")
save(fig, "fig09_behaviour_pareto")

# Fig 10 behaviour trace
r0, _ = s.behaviour(o["x"], o["P"], x, 0, None, False)
r1, _ = s.behaviour(o["x"], o["P"], x, 2, 0.15, True)
ideal = np.clip(x @ s.W.T, -s.R_MAX, s.R_MAX)
fig, ax = plt.subplots(figsize=(W1, 1.45))
ax.plot(tt, ideal[:, 0], "k", lw=0.7, label="Oracle $Wx$")
ax.plot(tt, r0[:, 0], color="#d0743c", lw=0.5, alpha=0.85, label="Baseline (λ=0)")
ax.plot(tt, r1[:, 0], color="#1b3f6b", lw=0.9, label="Selected (λ=2, τ=0.15)")
ax.set_xlim(20, 70); ax.set_xlabel("Time (s)"); ax.set_ylabel("Displayed valence $r_V$"); ax.legend(frameon=False, ncol=3, loc="upper center", bbox_to_anchor=(0.5, 1.3))
save(fig, "fig10_behaviour_trace")
print("done")