from __future__ import annotations

from pathlib import Path

import numpy as np
import torch

from ..mb_potential import simulate_overdamped_mb


BASIN_A = (-0.56, 1.44)
BASIN_B = (0.62, 0.03)
DEFAULT_DATA = Path(__file__).resolve().parents[2] / "data" / "mb_reactive"
K = 100
START_XY = BASIN_A


def generate(n_reps=12, steps=250_000, dt=2e-4, beta=0.06, mu=1.0, stride=5, seed=0):
    """Return a list of (T, 2) reactive MB trajectories (after striding).

    ``stride`` thins the saved frames (the integrator step dt is tiny for
    stability at high T; the saved frame spacing is dt*stride).
    """
    trajs = []
    for r in range(n_reps):
        x0, y0 = BASIN_A if r % 2 == 0 else BASIN_B
        xs, ys = simulate_overdamped_mb(
            n_steps=steps,
            dt=dt,
            beta=beta,
            mu=mu,
            x0=x0,
            y0=y0,
            noise=True,
            seed=seed + r,
            dtype=torch.float64,
        )
        xy = np.stack([xs, ys], axis=1)[::stride]
        if np.all(np.isfinite(xy)):
            trajs.append(xy.astype(np.float64))
    return trajs


def save(trajs, out_dir):
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    for f in out.glob("traj_*.npy"):
        f.unlink()
    for i, t in enumerate(trajs):
        np.save(out / f"traj_{i:04d}.npy", t)
    with open(out / "index.txt", "w") as fh:
        fh.write("\n".join(f"traj_{i:04d}.npy" for i in range(len(trajs))) + "\n")
    return out


def load_trajs(data_dir):
    files = sorted(Path(data_dir).glob("traj_*.npy"))
    if not files:
        print(f"  no reactive data in {data_dir} -- generating trajectories...")
        trajs = generate()
        save(trajs, data_dir)
        files = sorted(Path(data_dir).glob("traj_*.npy"))
    trajs = [np.load(f).astype(np.float64) for f in files]
    return trajs


def orient_endpoints(model, A, B):
    """Return (start_set, end_set) so start is the basin nearest START_XY."""
    xyA = model.active_to_xy(A).mean(axis=0)
    xyB = model.active_to_xy(B).mean(axis=0)
    dA = np.linalg.norm(xyA - np.array(START_XY))
    dB = np.linalg.norm(xyB - np.array(START_XY))
    return (A, B) if dA <= dB else (B, A)


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(description="Generate reactive MB trajectories")
    ap.add_argument("--out", default="data/mb_reactive")
    ap.add_argument("--n_reps", type=int, default=12)
    ap.add_argument("--steps", type=int, default=250_000)
    ap.add_argument("--dt", type=float, default=2e-4)
    ap.add_argument("--beta", type=float, default=0.06)
    ap.add_argument("--stride", type=int, default=5)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    trajs = generate(
        n_reps=args.n_reps,
        steps=args.steps,
        dt=args.dt,
        beta=args.beta,
        stride=args.stride,
        seed=args.seed,
    )
    out = save(trajs, args.out)
    nf = sum(len(t) for t in trajs)
    print(f"wrote {len(trajs)} trajectories ({nf} frames) -> {out}")