"""Exploration of the struck sheet (plate.py). Writes results/explore.npz; plots.py draws it.

    python experiments.py explore     # undamped: bending x dome x strike energy; energy above 1 / 4 kHz over time

Time is in seconds of sound with the lowest mode at F_LOW Hz, as in the web app.
"""

from __future__ import annotations

import os
import sys
from multiprocessing import Pool

os.environ.setdefault("OMP_NUM_THREADS", "1")

import numpy as np  # noqa: E402

from plate import Plate, run, strike  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
N, CI, CJ, RAD, F_LOW = 24, 8, 10, 2.0, 100.0


def _one(args):
    Kb, Hd, E = args
    pl = Plate(N=N, Kb=Kb, Hd=Hd)
    om = pl.omega()
    tau = om.min() / (2 * np.pi * F_LOW)
    dt = 0.6 / np.sqrt(om.max() ** 2 + 12 * E / N)
    rec = max(1, int(round(0.002 / tau / dt)))
    w0, v0 = strike(pl, E, CI, CJ, RAD, M=2, scale=[1.0, 1 + 1e-6])
    r = run(pl, w0, v0, 0.3 / tau, dt, record_every=rec)
    f = (om / om.min() * F_LOW).ravel()
    Ek = r["Ek"][:, 0].reshape(len(r["t"]), -1)
    return Kb, Hd, E, r["t"] * tau, Ek, f, r["E"][:, 0]


def explore(pool):
    tag = os.environ.get("TAG", "explore")
    if tag == "explore2":
        jobs = [(Kb, Hd, E) for Kb in (8.0, 25.0) for Hd in (0.0, 10.0) for E in (3000.0, 10000.0, 30000.0)]
    else:
        jobs = [(Kb, Hd, E) for Kb in (8.0, 25.0) for Hd in (0.0, 4.0) for E in (10.0, 100.0, 1000.0)]
    out = {}
    for Kb, Hd, E, t, Ek, f, Etot in pool.imap_unordered(_one, jobs):
        k = f"Kb{Kb:g}_Hd{Hd:g}_E{E:g}"
        out[k + "_t"], out[k + "_Ek"], out[k + "_f"], out[k + "_E"] = t, Ek, f, Etot
        sh = Ek / Ek.sum(1, keepdims=True)
        h1, h4 = sh[:, f > 1000].sum(1), sh[:, f > 4000].sum(1)
        i10, i50 = np.argmin(abs(t - 0.01)), np.argmin(abs(t - 0.05))
        print(f"  Kb {Kb:g} Hd {Hd:g} E {E:g}: top {f.max():.0f} Hz; above 1 kHz at 10/50/300 ms {h1[i10]:.2f} {h1[i50]:.2f} "
              f"{h1[-1]:.2f}; above 4 kHz {h4[i10]:.3f} {h4[i50]:.3f} {h4[-1]:.3f}; energy drift {abs(Etot[-1] / Etot[0] - 1):.1e}",
              flush=True)
    np.savez(os.path.join(RES, tag + ".npz"), **out)


if __name__ == "__main__":
    os.makedirs(RES, exist_ok=True)
    with Pool(int(os.environ.get("WORKERS", "12"))) as pool:
        explore(pool)