"""Plots for analysis/chaos-check.ts: python3 analysis/plot_chaos.py chaos-check.json out_dir"""

import json
import sys
from pathlib import Path

import matplotlib

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

d = json.load(open(sys.argv[1]))
out = Path(sys.argv[2] if len(sys.argv) > 2 else ".")
out.mkdir(parents=True, exist_ok=True)
t = np.array(d["time"]) * 1000  # ms
TMAX = 30


def sep_axes(ax, title):
    ax.set_yscale("log")
    ax.set_xlim(0, TMAX)
    ax.set_ylim(1e-12, 3e-2)
    ax.set_xlabel("time (ms)")
    ax.set_ylabel("RMS separation of the beads (m)")
    ax.set_title(title, fontsize=10)
    ax.grid(alpha=0.3)


# 1. step size
fig, ax = plt.subplots(figsize=(6.5, 4))
for s in d["stepSize"]:
    y = np.maximum(np.array(s["d"]), 1e-13)
    ax.plot(t, y, lw=1, label=f"cfl {s['cfl']} vs {s['against']}")
sep_axes(ax, "Same pluck (10 mm), time step halved: runs agree for longer")
ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(out / "step-size.png", dpi=150)

# 2. perturbations
fig, axs = plt.subplots(1, 2, figsize=(11, 4), sharey=True)
for ax, cfl in zip(axs, sorted({p["cfl"] for p in d["perturbation"]}, reverse=True)):
    for p in [p for p in d["perturbation"] if p["cfl"] == cfl]:
        y = np.maximum(np.array(p["d"]), 1e-13)
        g = p["growth"]
        lab = f"eps {p['eps']:g}" + (f", growth {g['lambda']:.0f}/s" if g else "")
        ax.plot(t, y, lw=1, label=lab)
    sep_axes(ax, f"Pluck amplitude times (1 + eps), cfl {cfl}")
    ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(out / "perturbation.png", dpi=150)

# 3. energy error vs time step
fig, ax = plt.subplots(figsize=(5, 4))
dt = np.array([r["dt"] for r in d["runs"]])
err = np.array([r["energyError"] for r in d["runs"]])
ax.loglog(dt, err, "o-", label="largest energy balance error / peak energy")
ax.loglog(dt, err[0] * (dt / dt[0]) ** 2, "k--", lw=0.8, label="slope 2")
ax.set_xlabel("time step (s)")
ax.set_ylabel("relative energy error")
ax.legend(fontsize=8)
ax.grid(alpha=0.3, which="both")
fig.tight_layout()
fig.savefig(out / "energy.png", dpi=150)

# 4. sound statistics: octave bands, step sizes vs perturbations
fig, ax = plt.subplots(figsize=(6.5, 4))
for r in d["runs"]:
    ax.plot([b["f"] for b in r["bands"]], [b["db"] for b in r["bands"]], "o-", color="C0", alpha=0.6, lw=1)
for p in d["perturbation"]:
    ax.plot([b["f"] for b in p["bands"]], [b["db"] for b in p["bands"]], "s--", color="C1", alpha=0.6, lw=1)
ax.plot([], [], "o-", color="C0", label="6 time steps")
ax.plot([], [], "s--", color="C1", label="8 tiny perturbations")
ax.set_xscale("log")
ax.set_xlabel("octave band centre (Hz)")
ax.set_ylabel("level, 1 s to 2 s (dB re loudest band)")
ax.set_title("Sound statistics: time step vs perturbation", fontsize=10)
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(out / "bands.png", dpi=150)

# 5. aliasing
fig, ax = plt.subplots(figsize=(6.5, 4))
for a in d["aliasing"]:
    ax.plot([b["f"] for b in a["bands"]], [b["db"] for b in a["bands"]], "o-",
            label=f"{a['amp'] * 1000:g} mm: {a['fractionAbove24k'] * 100:.1f}% of power above 24 kHz")
ax.axvline(24000, color="k", lw=0.8, ls="--")
ax.text(25000, -5, "24 kHz", fontsize=8)
ax.set_xscale("log")
ax.set_xlabel("octave band centre (Hz)")
ax.set_ylabel("pickup level (dB re loudest band)")
ax.set_title("Pickup at 768 kHz: content above the 48 kHz Nyquist limit", fontsize=10)
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(out / "aliasing.png", dpi=150)
print("wrote", out)