"""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

from model import spectral_entropy  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")


def scan():
    d = np.load(os.path.join(RES, "scan.npz"))
    energies = d["energies"]
    fig, ax = plt.subplots(1, 3, figsize=(17, 4.6))
    rows = []
    for mu, col in ((0.0, "C0"), (1.0, "C3")):
        lam = np.array([np.median(d[f"mu{mu:g}_E{i}_lyap"]) for i in range(len(energies))])
        nfin = np.array([np.median(d[f"mu{mu:g}_E{i}_neff"][-1]) for i in range(len(energies))])
        ax[0].loglog(energies, lam, "o-", color=col, label=f"mu = {mu:g}")
        ax[1].semilogx(energies, nfin, "o-", color=col, label=f"mu = {mu:g}")
        for i in range(0, len(energies), 3):
            ne = np.median(d[f"mu{mu:g}_E{i}_neff"], axis=1)
            t = 20.0 * (np.arange(len(ne)) + 1)          # records every 20 time units
            ax[2].semilogx(t, ne, color=col, alpha=0.3 + 0.7 * i / len(energies), lw=1)
        rows.append((mu, lam, nfin))
    ax[0].set(xlabel="strike energy E (N = 32, beta = 1)", ylabel="Lyapunov exponent (4000 time units)",
              title="Rate of separation")
    ax[1].set(xlabel="strike energy E", ylabel="effective fraction of modes at t = 4000",
              title="Equipartition (about 0.7 for random equipartition)")
    ax[2].set(xlabel="time", ylabel="effective fraction of modes", title="Spreading over time (darker: more energy)")
    for a in ax:
        a.grid(True, which="both", alpha=0.3)
    ax[0].legend()
    ax[1].legend()
    fig.tight_layout()
    fig.savefig(os.path.join(RES, "scan.png"), dpi=100)
    return energies, rows


def cascade():
    d = np.load(os.path.join(RES, "cascade.npz"))
    Es = (0.5, 2.0, 10.0)
    fig, ax = plt.subplots(2, 3, figsize=(17, 7.5), sharex=True)
    for i, mu in enumerate((0.0, 1.0)):
        for j, E in enumerate(Es):
            k = f"mu{mu:g}_E{E:g}"
            t, Ek, om = d[k + "_t"], d[k + "_Ek"], d[k + "_omega"]
            frac = Ek / Ek.sum(1, keepdims=True)
            im = ax[i, j].pcolormesh(t, np.arange(1, Ek.shape[1] + 1), np.log10(frac.T + 1e-12),
                                     vmin=-6, vmax=0, cmap="magma", shading="auto", rasterized=True)
            ax[i, j].set(title=f"mu = {mu:g}, E = {E:g}", ylabel="mode k" if j == 0 else None)
            if i == 1:
                ax[i, j].set_xlabel("time")
    fig.colorbar(im, ax=ax, label="log10 share of energy in mode k", shrink=0.8)
    fig.savefig(os.path.join(RES, "cascade.png"), dpi=100, bbox_inches="tight")


def damped(mu):
    d = np.load(os.path.join(RES, "damped.npz"))
    E0s = (0.5, 5.0, 30.0)
    fig, ax = plt.subplots(4, 3, figsize=(17, 14))
    summary = []
    for j, E0 in enumerate(E0s):
        k = f"mu{mu:g}_E{E0:g}"
        t, Ek, E, logsum = d[k + "_t"], d[k + "_Ek"], d[k + "_E"], d[k + "_logsum"]
        neff = spectral_entropy(Ek)                     # (T, M)
        ax[0, j].semilogy(t, E[:, 0], "k")
        ax[0, j].set(title=f"mu = {mu:g}, strike energy E0 = {E0:g}", ylabel="energy")
        ax[1, j].plot(t, neff, lw=0.6)
        ax[1, j].set(ylabel="effective fraction of modes\n(16 strikes differing by 1e-10)", ylim=(0, 1))
        frac = Ek[-1] / Ek[-1].sum(1, keepdims=True)      # (M, N)
        ks = np.arange(1, frac.shape[1] + 1)
        for m in range(frac.shape[0]):
            ax[2, j].semilogy(ks, frac[m], lw=0.7)
        ax[2, j].set(xlabel="mode k", ylabel="final share of energy", ylim=(1e-8, 1))
        w = max(1, int(round(100 / (t[1] - t[0]))))
        loc = (logsum[w:, 0] - logsum[:-w, 0]) / (t[w:] - t[:-w])
        ax[3, j].plot(t[w:], loc, "k", lw=0.8)
        ax[3, j].axhline(0, color="0.6", lw=0.5)
        ax[3, j].set(xlabel="time", ylabel="local Lyapunov exponent (100-unit windows)")
        top = frac.mean(0) > 1e-3
        spread = np.log10(frac[:, top]).std(0)
        summary.append((E0, float(np.median(spread)), float(spread.max()), int(top.sum())))
    fig.tight_layout()
    fig.savefig(os.path.join(RES, f"damped_mu{mu:g}.png"), dpi=95)
    return summary


if __name__ == "__main__":
    if os.path.exists(os.path.join(RES, "scan.npz")):
        energies, rows = scan()
        for mu, lam, nfin in rows:
            print(f"mu {mu:g}: " + "  ".join(f"E {E:.2g}: lam {l:.4f} n {n:.2f}" for E, l, n in zip(energies, lam, nfin)))
    if os.path.exists(os.path.join(RES, "cascade.npz")):
        cascade()
    if os.path.exists(os.path.join(RES, "damped.npz")):
        for mu in (0.0, 1.0):
          for E0, med, mx, n in damped(mu):
            print(f"damped mu {mu:g}, E0 {E0:g}: spread of final log10 mode shares across strikes over the {n} modes "
                  f"holding > 0.1%: median {med:.3f}, max {mx:.3f} (decades)")
    print("done")