# OID_analysis/li22b/collect_reference.py
"""Collect reference fields: q_in (empty channel) and q_blk (uncontrolled pinball).

Usage:
  conda run -n pycuda_3_10 python3 src/OID_analysis/li22b/collect_reference.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; 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")
STEPS = 200


def collect_empty_channel(device):
    """q_in: no cylinders, just uniform flow."""
    out_dir = os.path.join(OUT_BASE, "q_in")
    if os.path.isfile(os.path.join(out_dir, "fields.npz")):
        print("q_in: SKIP (exists)")
        return
    os.makedirs(out_dir, exist_ok=True)
    sim = Simulation(CFG_PATH, device_id=device)
    sim.initialize()
    sim.run(4000)
    ux_list, uy_list = [], []
    for _ in range(STEPS):
        sim.run(1)
        macro = sim.get_macroscopic()
        ux_list.append(macro["ux"].copy())
        uy_list.append(macro["uy"].copy())
    sim.close()
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.array(ux_list, dtype=np.float32),
                        uy=np.array(uy_list, dtype=np.float32))
    print(f"q_in: saved {STEPS} snaps")


def collect_uncontrolled_pinball(device):
    """q_blk: pinball with b=[0,0,0]."""
    out_dir = os.path.join(OUT_BASE, "q_blk")
    if os.path.isfile(os.path.join(out_dir, "fields.npz")):
        print("q_blk: SKIP (exists)")
        return
    os.makedirs(out_dir, exist_ok=True)
    sim = Simulation(CFG_PATH, device_id=device)
    sim.add_body("circle", center=(1000.0, CY, 0.0), radius=R_CYL)
    sim.add_body("circle", center=(1026.0, CY + 15.0, 0.0), radius=R_CYL)
    sim.add_body("circle", center=(1026.0, CY - 15.0, 0.0), radius=R_CYL)
    sim.initialize()
    sim.run(4000)
    ux_list, uy_list = [], []
    for _ in range(STEPS):
        sim.run(1)
        macro = sim.get_macroscopic()
        ux_list.append(macro["ux"].copy())
        uy_list.append(macro["uy"].copy())
    sim.close()
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.array(ux_list, dtype=np.float32),
                        uy=np.array(uy_list, dtype=np.float32))
    print(f"q_blk: saved {STEPS} snaps")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", type=int, default=2)
    args = ap.parse_args()
    collect_empty_channel(args.device)
    collect_uncontrolled_pinball(args.device)
    print("Done.")


if __name__ == "__main__":
    main()
