"""Where does the struck sheet become chaotic? Pairs of strikes differing by 1e-6 in strength, with the web app's
damping and time stepping; how different are they after a second?

    WORKERS=32 python chaos_scan.py      # results/chaos_scan.npz and a printed table

For each setting: bending Kb, dome Hd, strike energy E, ring time of the highest mode (lowest mode 6 s), 1.5 s of
sound with the lowest mode at 100 Hz. Reported: relative difference of the pickup velocity in 1.0-1.5 s, and the
RMS difference of the final energy share per log-spaced band (dB, bands holding > 0.1%).
"""

from __future__ import annotations

import itertools
import os
from multiprocessing import Pool

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

import numpy as np  # noqa: E402

from plate import Plate, strike  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
N, SR, F1, DUR, DELTA = 24, 44100, 100.0, 1.5, 1e-6


def one(args):
    Kb, Hd, E, thi = args
    pl0 = Plate(N=N, Kb=Kb, Hd=Hd)
    om = pl0.omega()
    om1, omN = om.min(), om.max()
    tau = om1 / (2 * np.pi * F1)
    M = N + 1
    s1, sN = 4 * np.sin(np.pi / (2 * M)) ** 2, 4 * np.sin(N * np.pi / (2 * M)) ** 2
    lam1, lamN = 2 * s1, 2 * sN
    sLo, sHi = np.log(1000) / 6.0 * tau, np.log(1000) / thi * tau
    c1 = max(0.0, 2 * (sHi - sLo) / (lamN - lam1))
    c0 = max(0.0, 2 * sLo - c1 * lam1)
    pl = Plate(N=N, Kb=Kb, Hd=Hd, c0=c0, c1=c1)
    dtOut = 1 / SR / tau
    nsub = int(np.ceil(dtOut / (1.2 / omN)))
    h = dtOut / nsub
    w, v = strike(pl, E, 9, 10, 1.6, M=2, scale=[1.0, 1 + DELTA])
    n = int(DUR * SR)
    sig = np.zeros((n, 2))
    f = pl.force(w)
    for i in range(n):
        for _ in range(nsub):
            v = v + 0.5 * h * f
            w = w + h * v
            f = pl.force(w)
            v = v + 0.5 * h * f
            v = v - h * (pl.c0 * v - pl.c1 * pl.lap(v))
        if not np.isfinite(v).all():
            return args, np.nan, np.nan, True
        sig[i] = v[:, 15, 17]
    a, b = sig[int(1.0 * SR):, 0], sig[int(1.0 * SR):, 1]
    wave = float(np.linalg.norm(a - b) / np.linalg.norm(a))
    Ek = pl.mode_energies(w, v).reshape(2, -1)
    fk = (om / om1 * F1).ravel()
    edges = np.geomspace(F1, fk.max() * 1.0001, 25)
    bands = np.stack([Ek[:, (fk >= lo) & (fk < hi)].sum(1) for lo, hi in zip(edges[:-1], edges[1:])], 1)
    bands /= bands.sum(1, keepdims=True)
    keep = bands.mean(0) > 1e-3
    spec = float(np.sqrt(np.mean((10 * np.log10(bands[0, keep] / bands[1, keep])) ** 2)))
    return args, wave, spec, False


if __name__ == "__main__":
    os.makedirs(RES, exist_ok=True)
    grid = list(itertools.product((8.0, 25.0), (0.0, 10.0), (1e4, 3e4, 1e5), (0.8, 2.0, 5.0)))
    rows = []
    with Pool(int(os.environ.get("WORKERS", "32"))) as pool:
        for args, wave, spec, blew in pool.imap_unordered(one, grid):
            rows.append((*args, wave, spec, blew))
            Kb, Hd, E, thi = args
            print(f"Kb {Kb:4g} dome {Hd:4g} E {E:7g} high ring {thi:3g} s: "
                  + ("BLEW UP" if blew else f"waveform diff 1-1.5 s {wave:.1e}, final band spectra differ {spec:.2f} dB"), flush=True)
    np.savez(os.path.join(RES, "chaos_scan.npz"), rows=np.array(rows, dtype=float))