"""Figure from results/explore.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")


def explore():
    d = np.load(os.path.join(RES, "explore.npz"))
    cases = [(0.0, 0.0), (4.0, 0.0), (4.0, 3.0), (4.0, 8.0)]
    Es = (2.0, 10.0, 50.0)
    fig, ax = plt.subplots(2, 4, figsize=(20, 8.5))
    for j, (Kb, Hd) in enumerate(cases):
        k = f"Kb{Kb:g}_Hd{Hd:g}_E50"
        t, Ek, f = d[k + "_t"], d[k + "_Ek"], d[k + "_f"]
        share = Ek / Ek.sum(1, keepdims=True)
        ax[0, j].pcolormesh(t * 1000, f, np.log10(share.T + 1e-12), vmin=-6, vmax=0, cmap="magma", shading="auto",
                            rasterized=True)
        ax[0, j].set(yscale="log", ylim=(f[0], f[-1]), title=f"bending {Kb:g}, dome {Hd:g}, E = 50",
                     xlabel="time (ms of sound)", ylabel="mode pitch (Hz)" if j == 0 else None)
        for E, col in zip(Es, ("C0", "C1", "C3")):
            k = f"Kb{Kb:g}_Hd{Hd:g}_E{E:g}"
            t, Ek, f = d[k + "_t"], d[k + "_Ek"], d[k + "_f"]
            share = Ek / Ek.sum(1, keepdims=True)
            ax[1, j].plot(t * 1000, share[:, f > 4000].sum(1), color=col, label=f"E = {E:g}, above 4 kHz")
            ax[1, j].plot(t * 1000, share[:, f > 1000].sum(1), color=col, lw=0.8, ls="--", label=f"E = {E:g}, above 1 kHz")
        ax[1, j].set(ylim=(0, 1), xlabel="time (ms of sound)", ylabel="share of energy" if j == 0 else None)
        ax[1, j].grid(alpha=0.3)
    ax[1, 0].legend(fontsize=7)
    fig.suptitle("Undamped, 64 beads, lowest mode at 100 Hz: energy share per mode (top, log color, E = 50) and above 1 / 4 kHz (bottom)")
    fig.tight_layout()
    fig.savefig(os.path.join(RES, "explore.png"), dpi=90)


if __name__ == "__main__":
    explore()
    print("wrote results/explore.png")