"""Exploration of the struck wire (wire.py). Writes results/<name>.npz; plots.py draws them.

    python experiments.py explore     # undamped: bending x dome x strike energy; high-band share over time

Time is reported 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 wire import Wire, run, strike  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
N, CENTER, WIDTH, F_LOW = 64, 14, 3.0, 100.0


def dt_for(w: Wire, E):
    om, _ = w.modes()
    return 0.4 / np.sqrt(om[-1] ** 2 + 12 * E)


def tau_for(w: Wire):
    om, _ = w.modes()
    return om[0] / (2 * np.pi * F_LOW)          # seconds of sound per time unit


def _explore_one(args):
    Kb, Hd, E = args
    w = Wire(N=N, Kb=Kb, Hd=Hd)
    tau = tau_for(w)
    dt = dt_for(w, E)
    t_end = 0.3 / tau
    rec = max(1, int(round(0.002 / tau / dt)))
    y0, v0 = strike(w, E, CENTER, WIDTH, M=2, scale=[1.0, 1 + 1e-10])
    r = run(w, y0, v0, t_end, dt, record_every=rec)
    f = r["omega"] / (2 * np.pi * tau)          # mode frequencies in Hz of sound
    return Kb, Hd, E, r["t"] * tau, r["Ek"][:, 0, :], f, r["E"][:, 0]


def explore(pool):
    jobs = [(Kb, Hd, E) for Kb in (0.0, 4.0) for Hd in (0.0, 3.0, 8.0) for E in (2.0, 10.0, 50.0)]
    out = {}
    for Kb, Hd, E, t, Ek, f, Etot in pool.imap_unordered(_explore_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
        share = Ek / Ek.sum(1, keepdims=True)
        hi1, hi4 = share[:, f > 1000].sum(1), share[:, f > 4000].sum(1)
        print(f"  Kb {Kb:g} Hd {Hd:g} E {E:g}: top {f[-1]:.0f} Hz; share above 1 kHz at 10/50/300 ms "
              f"{hi1[4]:.2f} {hi1[24]:.2f} {hi1[-1]:.2f}; above 4 kHz {hi4[4]:.3f} {hi4[24]:.3f} {hi4[-1]:.3f}; "
              f"energy drift {abs(Etot[-1] / Etot[0] - 1):.1e}", flush=True)
    np.savez(os.path.join(RES, "explore.npz"), **out)


if __name__ == "__main__":
    os.makedirs(RES, exist_ok=True)
    which = sys.argv[1:] or ["explore"]
    with Pool(int(os.environ.get("WORKERS", "18"))) as pool:
        for w in which:
            print(w, flush=True)
            dict(explore=explore)[w](pool)