"""One-shot Dütsch Re 100 / KC 5 prescribed in-line oscillation driver.

Protocol: docs/validation/duetsch_re100_kc5_protocol.md. Native dense channel
(zero-velocity inlet, pressure-NEQ outlet, bounce-back y walls), MRT, lab frame.
Writes raw per-step force/diagnostics and phase-bin profile samples to --out
(outside the checkout); analysis is in analyze.py. No production change.

    python run.py --D 40 --box 40 --cycles 11 --out /path/outside/checkout
"""
from __future__ import annotations

import argparse
import json
import math
import os
import tempfile

import numpy as np

import protocol as p


TEMPLATE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "src",
                        "CelerisLab", "configs", "config_lbm.json")


def config(D, box, T):
    """Shipped config_lbm.json (full parser schema) with only the case fields
    overridden: grid, still fluid (velocity 0), nu from Re, MRT, uniform inlet
    profile, bounce-back y walls. Inlet scheme and outlet mode stay the
    template's (zou_he_local, neq_extrap) and are recorded in case.json."""
    u = p.KC * D / T
    n = int(round(box * D))
    with open(TEMPLATE) as fh:
        cfg = json.load(fh)
    cfg["grid"].update(lattice_model="D2Q9", nx=n, ny=n, nz=1)
    cfg["physics"].update(data_type="FP32", viscosity=u * D / p.RE, velocity=0.0, rho=1.0)
    cfg["method"].update(collision="MRT", streaming="double_buffer", store_precision="FP32",
                         ddf_shifting=False, y_wall_bc="bounce_back")
    cfg["method"]["les"]["enabled"] = False
    cfg["method"]["inlet"]["profile"] = "uniform"
    return cfg, u, n


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--D", type=int, required=True)
    ap.add_argument("--box", type=float, default=40.0, help="box side in D")
    ap.add_argument("--T", type=int, default=None, help="period in steps (multiple of 360)")
    ap.add_argument("--cycles", type=int, default=p.LAST_SCORED, help="run to the end of this cycle")
    ap.add_argument("--out", required=True)
    ap.add_argument("--dump-fields", default="",
                    help="research capture: comma-separated paper cycles, e.g. 11,21,30,31")
    args = ap.parse_args()
    D, T = args.D, args.T or 126 * args.D
    if T % 360:
        raise SystemExit("T must be a multiple of 360 steps")
    if os.path.realpath(args.out).startswith(os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))):
        raise SystemExit("--out must be outside the checkout")
    cfg, u_max, n = config(D, args.box, T)
    A, R = p.amplitude(D), 0.5 * D
    xm = ym = 0.5 * (n - 1)
    os.makedirs(args.out, exist_ok=True)

    from CelerisLab.body.motion import Translation
    from CelerisLab.simulation import Simulation
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as fh:
        json.dump(cfg, fh)
    sim = Simulation(fh.name)
    try:
        body = sim.add_body("circle", (xm + A, ym), R,
                            translation=Translation((-2.0 * A, 0.0, 0.0, 0.0), 1.01 * u_max))
        sim.initialize()
        dump = {int(c) for c in args.dump_fields.split(",") if c}
        if dump:
            import capture
            if max(dump) > args.cycles or min(dump) <= p.DISCARD:
                raise SystemExit("--dump-fields cycles must lie in the scored range of this run")
            crop = capture.crop_box(xm, ym, D, n, (xm - A - R, xm + A + R, ym - R, ym + R))
            trend = []
        steps = int(round((0.25 + args.cycles) * T))
        force = np.zeros((steps, 2)); diag = []; samples = []
        for m in range(steps):
            x0, x1 = p.position(m, T, A), p.position(m + 1, T, A)
            sim.set_body(body, offset=(max(x1 - A, -2.0 * A), 0.0), vx=x1 - x0)
            sim.run(1)
            force[m] = sim.read_force(body)[:2]
            diag.append(sim.moving_diagnostics()[0])
            s = m + 1  # state after the step
            ph = p.bin_phase(s, T)
            if ph is not None and p.cycle(s, T) > p.DISCARD:
                fields = sim.get_macroscopic()
                fluid = (sim.get_flags().reshape(n, n) & 1) != 0
                ux, uy = (np.asarray(fields[k]).reshape(n, n) for k in ("ux", "uy"))
                if dump:  # research capture only; the default path is unchanged
                    rho = np.asarray(fields["rho"]).reshape(n, n)
                    trend.append(capture.trend_row(s, p.cycle(s, T), ph, rho, ux, uy, fluid, D, n, crop))
                    if p.cycle(s, T) in dump and p.phase(s, T) == ph:  # first step of the bin
                        x0c, x1c, y0c, y1c = crop
                        cy_, cx_ = slice(y0c, y1c + 1), slice(x0c, x1c + 1)
                        np.savez_compressed(
                            os.path.join(args.out, f"fields_c{p.cycle(s, T)}_p{int(ph)}.npz"),
                            s=s, cycle=p.cycle(s, T), phase=ph, crop=np.asarray(crop),
                            force=force[m], centre=np.asarray((xm + p.position(s, T, A), ym)),
                            ux=ux[cy_, cx_].astype(np.float32), uy=uy[cy_, cx_].astype(np.float32),
                            rho=rho[cy_, cx_].astype(np.float32), fluid=fluid[cy_, cx_].astype(np.uint8),
                            down_ux=capture.downsample(ux, fluid), down_uy=capture.downsample(uy, fluid),
                            down_rho=capture.downsample(rho, fluid))
                centre = (xm + p.position(s, T, A), ym)  # analytic centre after the step
                for xs in p.STATIONS:
                    for x2 in p.X2:
                        x, y = xm + xs * D, ym + x2 * D
                        u, wu, wsum = p.sample(ux, fluid, x, y, centre, R)
                        v, _, _ = p.sample(uy, fluid, x, y, centre, R)
                        samples.append((s, p.cycle(s, T), ph, xs, x2,
                                        math.nan if u is None else u, math.nan if v is None else v,
                                        wsum, json.dumps(wu)))
        if dump:
            np.save(os.path.join(args.out, "trend.npy"), np.asarray(trend, float))
        np.savez(os.path.join(args.out, "raw.npz"), force=force, diag=np.asarray(diag),
                 samples=np.asarray([r[:8] for r in samples], float))
        with open(os.path.join(args.out, "case.json"), "w") as fh:
            json.dump(dict(D=D, T=T, box=args.box, n=n, u_max=u_max, nu=cfg["physics"]["viscosity"],
                           tau=3 * cfg["physics"]["viscosity"] + 0.5, A=A, steps=steps,
                           cycles=args.cycles, F_ref=0.5 * u_max ** 2 * D,
                           inlet=cfg["method"]["inlet"], outlet=cfg["method"]["outlet"],
                           weights=[r[8] for r in samples]), fh)
    finally:
        sim.close()


if __name__ == "__main__":
    main()