"""Figures from results/*.npz (written by experiments.py). Run: python plots.py"""

from __future__ import annotations

import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
CHAOS = 0.02    # finite-time Lyapunov exponent above which an orbit counts as chaotic


def sections():
    d = np.load(os.path.join(RES, "sections.npz"))
    energies, w2s = d["energies"], d["w2s"]
    fig, ax = plt.subplots(len(w2s), len(energies), figsize=(3 * len(energies), 3.1 * len(w2s)))
    for i, w2 in enumerate(w2s):
        for j, E in enumerate(energies):
            k = f"w2_{w2:g}_E_{E:g}"
            lam, owner, pts = d[k + "_lyap"], d[k + "_owner"], d[k + "_pts"]
            a = ax[i, j]
            chaotic = lam[owner] > CHAOS
            reg = ~chaotic
            a.scatter(pts[reg, 0], pts[reg, 1], s=0.15, c=owner[reg] % 10, cmap="tab10", vmin=0, vmax=9, rasterized=True)
            a.scatter(pts[chaotic, 0], pts[chaotic, 1], s=0.15, c="k", alpha=0.5, rasterized=True)
            a.set_title(f"w2 = {w2:g}, E = {E:g}\n{np.mean(lam > CHAOS):.0%} of orbits chaotic", fontsize=9)
            a.set_xticks([])
            a.set_yticks([])
            if j == 0:
                a.set_ylabel("p_x (section y = 0, p_y > 0)")
            if i == len(w2s) - 1:
                a.set_xlabel("x")
    fig.suptitle("Poincare sections: colored curves are regular orbits (tori), black dots chaotic orbits", y=1.0)
    fig.tight_layout()
    fig.savefig(os.path.join(RES, "sections.png"), dpi=110)


def lyapunov():
    d = np.load(os.path.join(RES, "lyapunov.npz"))
    energies, w2s = d["energies"], d["w2s"]
    fig, ax = plt.subplots(1, 2, figsize=(12, 4.5))
    lines = []
    for w2 in w2s:
        lams = [d[f"w2_{w2:g}_E{i}"] for i in range(len(energies))]
        frac = [np.mean(l > CHAOS) for l in lams]
        med = [np.median(l[l > CHAOS]) if np.any(l > CHAOS) else np.nan for l in lams]
        ax[0].semilogx(energies, frac, "o-", label=f"w2 = {w2:g}")
        ax[1].loglog(energies, med, "o-", label=f"w2 = {w2:g}")
        lines.append((w2, energies, frac, med))
    ax[0].set(xlabel="energy E (beta = 1)", ylabel="fraction of orbits chaotic", title="Onset of chaos")
    ax[1].set(xlabel="energy E", ylabel="median Lyapunov exponent of chaotic orbits",
              title="Rate of separation in the chaotic regime")
    for a in ax:
        a.grid(True, which="both", alpha=0.3)
        a.legend()
    fig.tight_layout()
    fig.savefig(os.path.join(RES, "lyapunov.png"), dpi=110)
    return lines


def damped():
    d = np.load(os.path.join(RES, "damped.npz"))
    E0s = sorted({float(k.split("_")[1]) for k in d.files if k.startswith("E0_")})
    fig, ax = plt.subplots(3, len(E0s), figsize=(5.5 * len(E0s), 10), sharex=True)
    summary = []
    for j, E0 in enumerate(E0s):
        k = f"E0_{E0:g}"
        t, Ex, Ey, E, logsum = d[k + "_t"], d[k + "_Ex"], d[k + "_Ey"], d[k + "_E"], d[k + "_logsum"]
        frac = Ex / (Ex + Ey)
        ax[0, j].semilogy(t, E[:, 0], "k")
        ax[0, j].set(title=f"strike energy E0 = {E0:g}", ylabel="energy (member 0)")
        ax[1, j].plot(t, frac, lw=0.6)
        ax[1, j].set(ylabel="share of energy in the x oscillator\n(16 strikes differing by 1e-10)", ylim=(0, 1))
        # local Lyapunov exponent: slope of the log-growth sum over windows of 20 time units
        w = max(1, int(round(20 / (t[1] - t[0]))))
        loc = (logsum[w:, 0] - logsum[:-w, 0]) / (t[w:] - t[:-w])
        ax[2, j].plot(t[w:], loc, "k", lw=0.8)
        ax[2, j].axhline(0, color="0.6", lw=0.5)
        ax[2, j].set(xlabel="time", ylabel="local Lyapunov exponent (20-unit windows)")
        late = t > t[-1] * 0.8
        summary.append((E0, float(frac[late].mean(0).std()), float(frac[late].mean(0).min()), float(frac[late].mean(0).max())))
    fig.tight_layout()
    fig.savefig(os.path.join(RES, "damped.png"), dpi=110)
    return summary


if __name__ == "__main__":
    for name, f in (("sections", sections), ("lyapunov", lyapunov), ("damped", damped)):
        if os.path.exists(os.path.join(RES, name + ".npz")):
            out = f()
            if name == "lyapunov":
                for w2, energies, frac, med in out:
                    print(f"w2 = {w2:g}:", " ".join(f"E {E:.2g}: {fr:.2f}" for E, fr in zip(energies, frac)))
            if name == "damped":
                for E0, sd, lo, hi in out:
                    print(f"damped E0 {E0:g}: late share of energy in x across strikes: std {sd:.3f}, range {lo:.3f} to {hi:.3f}")
            print("wrote", os.path.join(RES, name + ".png"))