"""Modal series solution of string-pluck-v1 (see the appendix of the challenge README).

Usage: python modal_string.py <challenge.json> <out-dir> [N ...]

Writes modal-n<N>.wav for each truncation order N and prints a JSON list of per-solution
metadata (timing, parameters) to stdout.
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from common import hardware_string, load_challenge, write_wav  # noqa: E402
from fields import field_grid, write_fields  # noqa: E402


def modal_coefficients(p: dict, n_modes: int):
    n = np.arange(1, n_modes + 1)
    k = n * np.pi / p["L"]
    q = np.pi / p["w"]
    A = (2.0 * p["h"] / p["L"]) * np.sin(k * p["x_p"]) * np.sin(k * p["w"]) * q**2 / (k * (q**2 - k**2))
    alpha = p["sigma0"] + 0.5 * p["kappa"] * k**2
    omega2 = (p["c"] * k) ** 2 - alpha**2
    if np.any(omega2 <= 0):
        raise ValueError("an overdamped mode is within the requested truncation; not handled")
    omega = np.sqrt(omega2)
    return k, A, alpha, omega


def pickup_velocity(p: dict, n_modes: int, t: np.ndarray) -> np.ndarray:
    k, A, alpha, omega = modal_coefficients(p, n_modes)
    coeff = -A * (p["c"] * k) ** 2 / omega * np.sin(k * p["x_m"])  # per-mode amplitude of a_n'(t) sin(k_n x_m)
    s = np.zeros_like(t)
    block = 64  # modes per block, to bound memory for large N
    for i in range(0, n_modes, block):
        sl = slice(i, i + block)
        e = np.exp(-np.outer(t, alpha[sl]))
        s += (e * np.sin(np.outer(t, omega[sl]))) @ coeff[sl]
    return s


def displacement_field(p: dict, n_modes: int, x: np.ndarray, t: np.ndarray) -> np.ndarray:
    """u(x_i, t_j) as an array of shape (nt, nx)."""
    k, A, alpha, omega = modal_coefficients(p, n_modes)
    u = np.zeros((len(t), len(x)))
    block = 64
    for i in range(0, n_modes, block):
        sl = slice(i, i + block)
        e = np.exp(-np.outer(t, alpha[sl]))
        a = A[sl] * e * (np.cos(np.outer(t, omega[sl])) + (alpha[sl] / omega[sl]) * np.sin(np.outer(t, omega[sl])))
        u += a @ np.sin(np.outer(k[sl], x))
    return u


def main():
    challenge = load_challenge(Path(sys.argv[1]))
    out_dir = Path(sys.argv[2])
    args = sys.argv[3:]
    fields_for = set()
    if "--fields" in args:
        i = args.index("--fields")
        fields_for = {int(v) for v in args[i + 1].split(",")}
        args = args[:i] + args[i + 2:]
    orders = [int(a) for a in args] or [4, 8, 16, 32, 64, 256, 1024]
    p = challenge["parameters"]
    out = challenge["output"]
    fs = out["sample_rate_hz"]
    t = np.arange(out["num_samples"]) / fs
    meta = []
    for n in orders:
        pickup_velocity(p, min(n, 8), t[:16])  # warm up
        t0 = time.perf_counter()
        s = pickup_velocity(p, n, t)
        dt = time.perf_counter() - t0
        write_wav(out_dir / f"modal-n{n}.wav", s, challenge)
        entry = {
            "solution_id": f"modal-n{n}",
            "output": f"{out_dir.name}/modal-n{n}.wav",
            "method": "modal series, truncated",
            "parameters": {"n_modes": n},
            "timing_seconds": round(dt, 4),
            "hardware": hardware_string(),
            "peak_si": float(np.max(np.abs(s))),
        }
        if n in fields_for:
            x, tt = field_grid(challenge)
            write_fields(out_dir / f"modal-n{n}.h5", x, tt, displacement_field(p, n, x, tt), challenge, f"modal-n{n}")
            entry["fields"] = {"path": f"{out_dir.name}/modal-n{n}.h5", "format": "hdf5"}
        meta.append(entry)
        print(f"N={n:5d}  peak={np.max(np.abs(s)):.4f} m/s  t={dt:.3f}s", file=sys.stderr)
    print(json.dumps(meta, indent=1))


if __name__ == "__main__":
    main()