"""Figure from results/explore2.npz (written by experiments.py with TAG=explore2). 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")

if __name__ == "__main__":
    d = np.load(os.path.join(RES, "explore2.npz"))
    fig, ax = plt.subplots(2, 2, figsize=(14, 8))
    for j, Kb in enumerate((8.0, 25.0)):
        k = f"Kb{Kb:g}_Hd10_E{10000:g}"
        t, Ek, f = d[k + "_t"], d[k + "_Ek"], d[k + "_f"]
        order = np.argsort(f)
        sh = Ek / Ek.sum(1, keepdims=True)
        edges = np.geomspace(f.min(), f.max() * 1.0001, 37)
        B = np.stack([sh[:, (f >= a) & (f < b)].sum(1) for a, b in zip(edges[:-1], edges[1:])], 1)
        ax[0, j].pcolormesh(t * 1000, np.sqrt(edges[:-1] * edges[1:]), np.log10(B.T + 1e-12), vmin=-6, vmax=0, cmap="magma",
                            shading="auto", rasterized=True)
        ax[0, j].set(yscale="log", title=f"bending {Kb:g}, dome 10, E = 10000: energy share per band", xlabel="time (ms of sound)",
                     ylabel="pitch (Hz)")
        for E, col in zip((3000.0, 10000.0, 30000.0), ("C0", "C1", "C3")):
            k = f"Kb{Kb:g}_Hd10_E{E:g}"
            t, Ek, f = d[k + "_t"], d[k + "_Ek"], d[k + "_f"]
            sh = Ek / Ek.sum(1, keepdims=True)
            ax[1, j].plot(t * 1000, sh[:, f > 4000].sum(1), color=col, label=f"E = {E:g}, above 4 kHz")
            ax[1, j].plot(t * 1000, sh[:, 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", title=f"bending {Kb:g}, dome 10, undamped")
        ax[1, j].grid(alpha=0.3)
    ax[1, 0].legend(fontsize=7)
    fig.tight_layout()
    fig.savefig(os.path.join(RES, "explore.png"), dpi=90)
    print("wrote results/explore.png")