"""Experiments with the struck FPUT-beta chain. Each writes results/<name>.npz; plots.py draws them.

    python experiments.py scan        # Lyapunov exponent and equipartition vs strike energy (undamped)
    python experiments.py cascade     # mode energies over time at three strike energies (undamped)
    python experiments.py damped      # damped strikes: ensembles of 16 nearby strikes, frozen outcome
    python experiments.py all

N = 32, beta = 1, strike centered on mass 8 with width 2 masses. mu = 0 (harmonic at low k) or mu = 1
(on-site springs: an inharmonic spectrum, closer to the plate).
"""

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 model import Chain, run, spectral_entropy, strike  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
N, CENTER, WIDTH = 32, 8, 2.0


def dt_for(E, mu):
    return min(0.1, 0.25 / np.sqrt(4 + mu**2 + 12 * E))


def _scan_one(args):
    mu, E = args
    ch = Chain(N=N, mu=mu)
    M = 6
    u0, v0 = strike(ch, E, CENTER, WIDTH, M=M, scale=1 + 1e-6 * np.arange(M))
    dt = dt_for(E, mu)
    r = run(ch, u0, v0, 4000.0, dt, lyap=True, record_every=int(round(20 / dt)))
    neff = spectral_entropy(r["Ek"])           # (T, M)
    return mu, E, r["lyap"], r["t"], neff


def scan(pool):
    energies = np.geomspace(0.05, 50, 16)
    out = {}
    for mu, E, lam, t, neff in pool.imap_unordered(_scan_one, [(mu, E) for mu in (0.0, 1.0) for E in energies]):
        i = int(np.argmin(abs(energies - E)))
        out[f"mu{mu:g}_E{i}_lyap"], out[f"mu{mu:g}_E{i}_neff"], out[f"mu{mu:g}_t"] = lam, neff, t
        print(f"  scan mu {mu:g} E {E:.3g}: lambda {np.median(lam):.4f}, n_eff/N final {np.median(neff[-1]):.2f}", flush=True)
    np.savez(os.path.join(RES, "scan.npz"), energies=energies, **out)


def _cascade_one(args):
    mu, E = args
    ch = Chain(N=N, mu=mu)
    u0, v0 = strike(ch, E, CENTER, WIDTH)
    dt = dt_for(E, mu)
    r = run(ch, u0, v0, 3000.0, dt, record_every=int(round(2 / dt)))
    return mu, E, r["t"], r["Ek"][:, 0, :], ch.omega


def cascade(pool):
    out = {}
    jobs = [(mu, E) for mu in (0.0, 1.0) for E in (0.5, 2.0, 10.0)]
    for mu, E, t, Ek, om in pool.imap_unordered(_cascade_one, jobs):
        k = f"mu{mu:g}_E{E:g}"
        out[k + "_t"], out[k + "_Ek"], out[k + "_omega"] = t, Ek, om
        print(f"  cascade mu {mu:g} E {E:g}", flush=True)
    np.savez(os.path.join(RES, "cascade.npz"), **out)


def damping_for(ch, sigma_low, sigma_high):
    """c0, c1 so that the lowest mode decays (amplitude) at sigma_low and the highest at sigma_high."""
    om = ch.omega
    a1, aN = om[0] ** 2 - ch.mu**2, om[-1] ** 2 - ch.mu**2
    c1 = 2 * (sigma_high - sigma_low) / (aN - a1)
    c0 = 2 * sigma_low - c1 * a1
    return c0, c1


def _damped_one(args):
    mu, E0 = args
    base = Chain(N=N, mu=mu)
    c0, c1 = damping_for(base, 4e-4, 3e-3)
    ch = Chain(N=N, mu=mu, c0=c0, c1=c1)
    M = 16
    u0, v0 = strike(ch, E0, CENTER, WIDTH, M=M, scale=1 + 1e-10 * np.arange(M))
    dt = dt_for(E0, mu)
    r = run(ch, u0, v0, 8000.0, dt, lyap=True, record_every=int(round(5 / dt)))
    return mu, E0, r["t"], r["Ek"], r["E"], r["logsum"], c0, c1


def damped(pool):
    out = {}
    jobs = [(mu, E0) for mu in (0.0, 1.0) for E0 in (0.5, 5.0, 30.0)]
    for mu, E0, t, Ek, E, logsum, c0, c1 in pool.imap_unordered(_damped_one, jobs):
        k = f"mu{mu:g}_E{E0:g}"
        out[k + "_t"], out[k + "_Ek"], out[k + "_E"], out[k + "_logsum"] = t, Ek, E, logsum
        out[k + "_c"] = np.array([c0, c1])
        print(f"  damped mu {mu:g} E0 {E0:g}", flush=True)
    np.savez(os.path.join(RES, "damped.npz"), **out)


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