"""One-shot predeclared VIV run (docs/validation/viv_re150_predeclare.md).

Re 150, U* = 5, m/(rho D^2) = 2, zeta = 0, D = 40 on a nominal 75 D x 60 D
lattice (uniform inlet, NEQ outlet, free-slip y walls). The cylinder is at
(30 D, (ny - 1)/2); y-wall planes are .5 and ny - 1.5 in lattice coordinates.
Structural RHS = CME + Wen (diagnostics: raw CME, node terms).
Research driver only: Simulation API, no production change.

    PYTHONPATH=src python research/viv_re150/run.py --out /path/outside/checkout
"""
from __future__ import annotations

import argparse
import json
import math
import os
import tempfile

import numpy as np

import viv

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


def config(D=D, U=U):
    nx, ny = 75 * D, 60 * D
    with open(TEMPLATE) as fh:
        cfg = json.load(fh)
    cfg["grid"].update(lattice_model="D2Q9", nx=nx, ny=ny, nz=1)
    cfg["physics"].update(data_type="FP32", viscosity=U * D / viv.RE, velocity=U, rho=1.0)
    cfg["method"].update(collision="MRT", streaming="double_buffer", store_precision="FP32",
                         ddf_shifting=False, y_wall_bc="free_slip")
    cfg["method"]["les"]["enabled"] = False
    cfg["method"]["inlet"]["profile"] = "uniform"
    return cfg, nx, ny


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--D", type=int, default=D)
    ap.add_argument("--smoke", action="store_true",
                    help="plumbing smoke only: exactly --steps steps (<= 64), never scored")
    ap.add_argument("--steps", type=int, default=0)
    args = ap.parse_args()
    if args.smoke != bool(args.steps) or args.steps > 64 or args.steps < 0:
        raise SystemExit("--smoke requires --steps N with 1 <= N <= 64 (and only with --smoke)")
    if args.D != D and not args.smoke:
        raise SystemExit("the physical case is predeclared at D = 40")
    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")
    os.makedirs(args.out, exist_ok=True)
    for name in ("report.json", "steps.npy"):  # never overwrite earlier evidence
        if os.path.exists(os.path.join(args.out, name)):
            raise SystemExit(f"{args.out} already holds {name}; use a fresh output directory")
    Dn = args.D
    cfg, nx, ny = config(Dn)
    s, cyc = viv.Structure(Dn, U), viv.Cycles(Dn)
    T_n = 1.0 / s.f_n
    step_cap, n_seed = int(round(viv.STEP_CAP_TN * T_n)), int(round(viv.SEED_TN * T_n))
    quiet = int(round(5 * T_n))
    if args.smoke:
        step_cap = args.steps
    centre = (30.0 * Dn, 0.5 * (ny - 1))

    from CelerisLab.body.motion import MovingTables, Translation
    col = {k: MovingTables.DIAG.index(k) for k in ("cme_y", "wen_y", "cover_jy", "uncover_jy", "no_donor")}
    from CelerisLab.simulation import Simulation
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as fh:
        json.dump(cfg, fh)
    sim = Simulation(fh.name)
    rows, status, seed_step, f_last = [], "running", None, 0.0
    try:
        body = sim.add_body("circle", centre, 0.5 * Dn,
                            translation=Translation((0.0, 0.0, -1.5 * Dn, 1.5 * Dn), 0.5))
        sim.initialize()
        for n in range(step_cap):
            # rows[k] holds the state y_{k+1}; the declared window is y_t, t in [n_seed - 5 T_n, n_seed)
            if not args.smoke and n == n_seed and \
                    max(abs(r[1]) for r in rows[n_seed - quiet - 1:n_seed - 1]) < viv.ELIGIBLE * Dn:
                s.v += s.seed_velocity(Dn)         # after the quiet check, before the update
                seed_step = n
                cyc.discard_current()
                cyc.reset_proof()
            y1, v1 = s.candidate(f_last)
            if not (math.isfinite(y1) and math.isfinite(v1)) or abs(y1) > viv.Y_LIMIT * Dn:
                status = "fail_bound_or_nonfinite"
                break
            try:
                sim.set_body(body, offset=(0.0, y1), vy=y1 - s.y)
            except ValueError as exc:              # input rejection: stop, no retry
                status = f"rejected: {exc}"
                break
            sim.run(1)                             # device failures raise and mark the Simulation failed
            d = sim.moving_diagnostics()[0]
            f_obs = sim.read_force(body)[1]
            f_rhs = float(d[col["cme_y"]] + d[col["wen_y"]])   # CME + Wen (y)
            row = (n, y1, v1, f_rhs, float(d[col["cme_y"]]), float(d[col["wen_y"]]), float(f_obs),
                   float(d[col["cover_jy"]] - d[col["uncover_jy"]]), float(d[col["no_donor"]]))
            if not viv.finite(row):                # every recorded output, not only the RHS
                status = "fail_nonfinite_output"
                break
            s.y, s.v, f_last = y1, v1, f_rhs       # commit only after a successful step
            rows.append(row)
            if args.smoke:
                continue                           # no settle / score in the plumbing smoke
            if cyc.update(s.y) == "scored":
                status = "scored"
                break
            if cyc.counted >= viv.CYCLE_CAP:
                status = "timeout_response_cycles"
                break
        else:
            status = ("smoke_ok" if all(r[8] == 0.0 for r in rows) else "smoke_no_donor") \
                if args.smoke else "timeout_native_steps"
    except BaseException as exc:
        status = f"error: {exc!r}"
        raise
    finally:
        np.save(os.path.join(args.out, "steps.npy"), np.asarray(rows, float))
        report = dict(status=status, smoke=args.smoke, D=Dn, steps=len(rows), t_U_over_D=len(rows) * U / Dn, T_n=T_n,
                      seed_step=seed_step, counted_cycles=cyc.counted, cycles=cyc.log,
                      columns=["n", "y", "v", "F_rhs_cme_wen", "F_cme", "F_wen", "F_obs",
                               "node_cover_minus_uncover_y", "no_donor"])
        if status == "scored" and not args.smoke:
            report["score"] = viv.score(cyc.window)
        with open(os.path.join(args.out, "report.json"), "w") as fh:
            json.dump(report, fh, indent=1)
        sim.close()


if __name__ == "__main__":
    main()