# OID_analysis/li22b/collect_openloop_db.py
"""Generate Li22b open-loop database: 50 steady control commands on fluidic pinball.

Uses new CelerisLab Simulation API, 2000x600 uniform-inlet config (no disturbance cylinder).
Saves ux/uy fields, 9 sensors (3x3 grid), and forces per command.

Usage:
  conda run -n pycuda_3_10 python3 src/OID_analysis/li22b/collect_openloop_db.py --device 2
"""
import os, sys, json, time, argparse
import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

from CelerisLab import Simulation

U0 = 0.01; L0 = 20.0; R_CYL = 10.0
NX, NY = 2000, 600
CY = (NY - 1) / 2.0
CFG_PATH = os.path.join(_REPO, "configs", "config_lbm_karman_2000x600.json")
OUT_BASE = os.path.join(os.path.dirname(__file__), "..", "data_li22b")
WARMUP_STEPS = 4000
RECORD_STEPS = 200
DEVICE = 2

PINBALL_FRONT_X = 1000.0
PINBALL_REAR_X  = 1026.0
PINBALL_Y_SPAN  = 15.0
SENSOR_XS = [PINBALL_REAR_X + dx for dx in [100.0, 130.0, 160.0]]
SENSOR_YS = [CY + 25.0, CY, CY - 25.0]


def latin_hypercube(n, d=3, low=-2.0, high=2.0, seed=42):
    rng = np.random.RandomState(seed)
    samples = np.zeros((n, d))
    for j in range(d):
        perm = rng.permutation(n)
        samples[:, j] = low + (high - low) * (perm + rng.rand(n)) / n
    return samples


def run_one(cmd_id, b):
    out_dir = os.path.join(OUT_BASE, f"{cmd_id:03d}")
    fields_fp = os.path.join(out_dir, "fields.npz")
    if os.path.isfile(fields_fp):
        print(f"  [{cmd_id:03d}] SKIP (exists)")
        return
    os.makedirs(out_dir, exist_ok=True)

    sim = Simulation(CFG_PATH, device_id=DEVICE)
    sim.add_body("circle", center=(PINBALL_FRONT_X, CY, 0.0), radius=R_CYL)
    sim.add_body("circle", center=(PINBALL_REAR_X,  CY + PINBALL_Y_SPAN, 0.0), radius=R_CYL)
    sim.add_body("circle", center=(PINBALL_REAR_X,  CY - PINBALL_Y_SPAN, 0.0), radius=R_CYL)
    for sx in SENSOR_XS:
        for sy in SENSOR_YS:
            sim.add_body("sensor", center=(sx, sy, 0.0), radius=5.0)
    sim.initialize()

    for i in range(3):
        sim.set_body(i, omega=float(b[i]) * U0 / R_CYL)

    sim.run(WARMUP_STEPS)

    ux_list, uy_list = [], []
    sens_list, force_list = [], []
    for step in range(RECORD_STEPS):
        sim.run(1)
        macro = sim.get_macroscopic()
        ux_list.append(macro["ux"].copy())
        uy_list.append(macro["uy"].copy())
        ss, fs = [], []
        for si in range(9):
            sv = sim.read_sensor(si, normalize=False)
            ss.extend([float(sv[0]), float(sv[1])])
        sens_list.append(ss)
        for fi in range(3):
            fv = sim.read_force(fi)
            fs.extend([float(fv[0]), float(fv[1])])
        force_list.append(fs)
        if (step + 1) % 100 == 0:
            print(f"    {step+1}/{RECORD_STEPS}")

    sim.close()
    np.savez_compressed(fields_fp,
                        ux=np.array(ux_list, dtype=np.float32),
                        uy=np.array(uy_list, dtype=np.float32))
    np.savez_compressed(os.path.join(out_dir, "sensors.npz"),
                        sensors=np.array(sens_list, dtype=np.float32))
    np.savez_compressed(os.path.join(out_dir, "forces.npz"),
                        forces=np.array(force_list, dtype=np.float32))
    with open(os.path.join(out_dir, "config.json"), "w") as f:
        json.dump({"b": b.tolist(), "nu": 0.004, "U0": U0,
                   "Re_D": 50, "Re_code": 100}, f, indent=2)
    print(f"  [{cmd_id:03d}] b={b[0]:+.3f},{b[1]:+.3f},{b[2]:+.3f}  saved")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--num", type=int, default=10)
    ap.add_argument("--start", type=int, default=0)
    args = ap.parse_args()
    global DEVICE; DEVICE = args.device

    b_all = latin_hypercube(args.num, seed=args.seed)
    print(f"{args.num} commands, device={DEVICE}, seed={args.seed}")
    print(f"Output: {OUT_BASE}")
    t0 = time.time()
    for i in range(args.start, args.num):
        run_one(i, b_all[i])
        print(f"  [ETA: {(time.time()-t0)*(args.num-i-1)/((i-args.start+1)*60):.0f} min]")
    print(f"Done in {(time.time()-t0)/60:.1f} min")


if __name__ == "__main__":
    main()
