"""Experiments with the two quartic-coupled oscillators. Each writes results/<name>.npz; plots.py draws them.

    python experiments.py sections     # Poincare sections at several energies (w2 = 1 and 1.618)
    python experiments.py lyapunov     # finite-time largest Lyapunov exponent vs energy
    python experiments.py damped       # a damped 'strike' and an ensemble of nearby strikes
    python experiments.py all          # all three, in parallel processes

beta = 1 throughout, so the energy E is the single parameter (beta E).
"""

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 System, run, section_initial_conditions  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
W2S = (1.0, 1.618)


def dt_for(E):
    """Step small enough for the stiffening at energy E (the effective frequency grows like sqrt(E))."""
    return min(0.01, 0.1 / np.sqrt(1 + 4 * E))


def _sections_one(args):
    w2, E = args
    sys_ = System(1.0, w2, 1.0)
    s0 = section_initial_conditions(sys_, E, 9)
    r = run(sys_, s0, 3000.0, dt=dt_for(E), lyap=True, section=True)
    pts = r["section"]
    return w2, E, r["lyap"], np.concatenate([np.full(len(p), i) for i, p in enumerate(pts)]), \
        np.concatenate([p for p in pts if len(p)] or [np.zeros((0, 2))])


def sections(pool):
    energies = (0.25, 0.5, 1.0, 1.5, 2.0, 4.0, 8.0)
    out = {}
    for w2, E, lam, owner, pts in pool.imap_unordered(_sections_one, [(w2, E) for w2 in W2S for E in energies]):
        key = f"w2_{w2:g}_E_{E:g}"
        out[key + "_lyap"], out[key + "_owner"], out[key + "_pts"] = lam, owner, pts
        print(f"  sections w2 {w2:g} E {E:g}: {len(lam)} orbits", flush=True)
    np.savez(os.path.join(RES, "sections.npz"), energies=np.array(energies), w2s=np.array(W2S), **out)


def _lyap_one(args):
    w2, E = args
    sys_ = System(1.0, w2, 1.0)
    s0 = section_initial_conditions(sys_, E, 16)
    r = run(sys_, s0, 2000.0, dt=dt_for(E), lyap=True)
    return w2, E, r["lyap"]


def lyapunov(pool):
    energies = np.geomspace(0.1, 100, 25)
    lam = {}
    for w2, E, l in pool.imap_unordered(_lyap_one, [(w2, E) for w2 in W2S for E in energies]):
        lam[(w2, E)] = l
        print(f"  lyapunov w2 {w2:g} E {E:.3g}: median {np.median(l):.3f}", flush=True)
    np.savez(os.path.join(RES, "lyapunov.npz"), energies=energies, w2s=np.array(W2S),
             **{f"w2_{w2:g}_E{i}": lam[(w2, E)] for w2 in W2S for i, E in enumerate(energies)})


def _damped_one(args):
    """A 'strike': both oscillators start at rest at the origin with momentum sqrt(2 E0) in direction a,
    then decay with p -> p exp(-gamma dt). An ensemble perturbs a by i * 1e-10. Records the mode energies
    and the running Lyapunov log-sum of member 0."""
    E0, gamma, n_members = args
    sys_ = System(1.0, 1.618, 1.0)
    a = 0.7 + 1e-10 * np.arange(n_members)
    p = np.sqrt(2 * E0)
    s0 = np.array([np.zeros(n_members), np.zeros(n_members), p * np.cos(a), p * np.sin(a)])
    r = run(sys_, s0, 1500.0, dt=dt_for(E0), gamma=gamma, lyap=True, renorm_every=10,
            record_every=int(round(0.5 / dt_for(E0))))
    Ex, Ey = sys_.mode_energies(np.moveaxis(r["states"], 1, 0))
    return E0, r["t"], Ex, Ey, sys_.energy(np.moveaxis(r["states"], 1, 0)), r["logsum"]


def damped(pool):
    gamma = 0.005
    out = {}
    for E0, t, Ex, Ey, Etot, logsum in pool.imap_unordered(_damped_one, [(E0, gamma, 16) for E0 in (0.5, 3.0, 20.0)]):
        k = f"E0_{E0:g}"
        out[k + "_t"], out[k + "_Ex"], out[k + "_Ey"], out[k + "_E"], out[k + "_logsum"] = t, Ex, Ey, Etot, logsum
        print(f"  damped E0 {E0:g} done", flush=True)
    np.savez(os.path.join(RES, "damped.npz"), gamma=gamma, **out)


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