from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np

from example_common import make_params
from lda.ks import integrate_steps, random_initial_state
from lda.utils import ensure_dir


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate a dispersive-KS truth trajectory.")
    parser.add_argument("--n", type=int, default=512)
    parser.add_argument("--length", type=float, default=None)
    parser.add_argument("--dt", type=float, default=0.05)
    parser.add_argument("--spinup-steps", type=int, default=2000)
    parser.add_argument("--save-steps", type=int, default=800)
    parser.add_argument("--save-every", type=int, default=20)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--output", type=Path, default=Path("example_outputs/truth_run.npz"))
    args = parser.parse_args()

    params = make_params(args.n, args.length)
    rng = np.random.default_rng(args.seed)
    state = random_initial_state(params, rng, std=0.5, max_initial_mode=min(64, params.n // 4))
    state = integrate_steps(state, args.spinup_steps, args.dt, params)

    states = [state.copy()]
    times = [args.spinup_steps * args.dt]
    for step in range(args.save_every, args.save_steps + 1, args.save_every):
        state = integrate_steps(state, args.save_every, args.dt, params)
        states.append(state.copy())
        times.append((args.spinup_steps + step) * args.dt)

    ensure_dir(args.output.parent)
    np.savez_compressed(
        args.output,
        truth=np.asarray(states),
        time=np.asarray(times),
        n=args.n,
        length=params.length,
        dt=args.dt,
        delta3=params.third_order_dispersion,
        delta5=params.fifth_order_dispersion,
    )
    print(f"saved {args.output}")


if __name__ == "__main__":
    main()
