"""Finite-difference solution of string-pluck-v1, second order in space and time.

With D the standard three-point second difference (zero at the fixed ends) and u^n the grid
values at step n, the scheme is

  (u^{n+1} - 2u^n + u^{n-1})/dt^2 = c^2 D u^n / dx^2
                                    - sigma0 (u^{n+1} - u^{n-1}) / dt
                                    + kappa D (u^{n+1} - u^{n-1}) / (2 dt dx^2).

The Kelvin-Voigt term is centered in time, so the scheme is implicit in the damping and a
tridiagonal system is solved at every step. (An earlier explicit variant with a backward
difference for that term is only first order in time; its error cancelled against the
second-order dispersion error near m = 8 and the sequence was not monotone.) Stability
requires the Courant number c dt/dx <= 1.

The time step is dt = 1/(fs*m) for an integer oversampling factor m, and the grid spacing is
dx = c dt / lam with lam = 0.9, so the scheme is deliberately dispersive. The pickup velocity is
(u^{n+1} - u^{n-1})/(2 dt) linearly interpolated to x_m, taken every m-th step.

Usage: python fd_string.py <challenge.json> <out-dir> [m ...]
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np
from scipy.linalg import solve_banded

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

LAM = 0.9


def solve(p: dict, fs: int, num_samples: int, m: int, lam: float = LAM, field_x=None, field_t=None):
    """Returns (pickup samples, info, fields) where fields is u on (field_t, field_x) or None.

    Field frames are taken at the step nearest each requested time and linearly interpolated in x."""
    dt = 1.0 / (fs * m)
    n_int = int(np.floor(p["L"] * lam / (p["c"] * dt)))  # intervals; floor keeps dx >= c dt / lam
    dx = p["L"] / n_int
    x = np.linspace(0.0, p["L"], n_int + 1)
    lam2 = (p["c"] * dt / dx) ** 2
    s = p["sigma0"] * dt
    beta = p["kappa"] * dt / (2.0 * dx**2)

    def D(u):
        out = np.zeros_like(u)
        out[1:-1] = u[:-2] - 2.0 * u[1:-1] + u[2:]
        return out

    # Left-hand operator on interior points: (1 + s) I - beta D, tridiagonal, in solve_banded form.
    n_in = n_int - 1
    ab = np.zeros((3, n_in))
    ab[0, 1:] = -beta
    ab[1, :] = (1.0 + s) + 2.0 * beta
    ab[2, :-1] = -beta

    u_prev = initial_shape(x, p)
    u_cur = u_prev + 0.5 * dt**2 * (p["c"] ** 2) * D(u_prev) / dx**2  # zero initial velocity, second-order start
    xi = p["x_m"] / dx
    i0 = int(np.floor(xi))
    w1 = xi - i0
    out = np.zeros(num_samples)
    total_steps = num_samples * m
    u_next = np.zeros_like(u_cur)
    fields = None
    frame_steps = {}
    if field_x is not None:
        fields = np.zeros((len(field_t), len(field_x)))
        for j, tj in enumerate(field_t):
            frame_steps.setdefault(int(round(tj / dt)), []).append(j)
        for j in frame_steps.get(0, []):
            fields[j] = np.interp(field_x, x, u_prev)
        for j in frame_steps.get(1, []):
            fields[j] = np.interp(field_x, x, u_cur)
    for n in range(1, total_steps + 1):
        rhs = 2.0 * u_cur + lam2 * D(u_cur) - (1.0 - s) * u_prev - beta * D(u_prev)
        u_next[1:-1] = solve_banded((1, 1), ab, rhs[1:-1], overwrite_b=True, check_finite=False)
        if n % m == 0:
            k = n // m
            if k < num_samples:
                v = (u_next - u_prev) / (2.0 * dt)
                out[k] = (1.0 - w1) * v[i0] + w1 * v[i0 + 1]
        for j in frame_steps.get(n + 1, []):
            fields[j] = np.interp(field_x, x, u_next)
        u_prev, u_cur, u_next = u_cur, u_next, u_prev
    return out, {"m": m, "dt": dt, "dx": dx, "n_points": n_int + 1, "courant": p["c"] * dt / dx}, fields


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:]
    ms = [int(a) for a in args] or [1, 2, 4, 8, 16]
    p = challenge["parameters"]
    out = challenge["output"]
    meta = []
    for m in ms:
        t0 = time.perf_counter()
        sig, info, _ = solve(p, out["sample_rate_hz"], out["num_samples"], m)
        dt = time.perf_counter() - t0
        sid = f"fd-m{m}"
        write_wav(out_dir / f"{sid}.wav", sig, challenge)
        entry = {
            "solution_id": sid,
            "output": f"{out_dir.name}/{sid}.wav",
            "method": "finite differences, second order, implicit damping",
            "parameters": {"oversampling": m, "n_points": info["n_points"], "dx_m": info["dx"], "dt_s": info["dt"], "courant": round(info["courant"], 6)},
            "timing_seconds": round(dt, 3),
            "hardware": hardware_string(),
            "peak_si": float(np.max(np.abs(sig))),
        }
        if m in fields_for:
            x, tt = field_grid(challenge)
            _, _, fld = solve(p, out["sample_rate_hz"], out["num_samples"], m, field_x=x, field_t=tt)
            write_fields(out_dir / f"{sid}.h5", x, tt, fld, challenge, sid)
            entry["fields"] = {"path": f"{out_dir.name}/{sid}.h5", "format": "hdf5"}
        meta.append(entry)
        print(f"m={m}  points={info['n_points']}  peak={np.max(np.abs(sig)):.4f} m/s  t={dt:.2f}s", file=sys.stderr)
    print(json.dumps(meta, indent=1))


if __name__ == "__main__":
    main()