"""Collect empty channel reference flow (parabolic profile, no bodies).

This provides the "target" reference for steady cloak metrics —
the clean channel flow that the cloak should restore.

Usage:
    conda run -n pycuda_3_10 python scripts/collect_empty_channel.py --device 2

Output: data/target_channel/target_channel/
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

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)
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from LegacyCelerisLab import FlowField

from CCD_analysis.configs import get_scene, data_dir_for_scene, LEGACY_CFG_DIR
from CCD_analysis.utils.cfd_interface import (
    load_legacy_configs, get_velocity_field, save_vorticity_png,
    vorticity_from_ddf,
)


def collect():
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", type=int, default=2)
    args = ap.parse_args()

    cfg = get_scene("target_channel")
    out_dir = data_dir_for_scene("target_channel")
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))

    ff = FlowField(field_cfg, cuda_cfg, device_id=args.device)
    from CCD_analysis.configs import L0, CENTER_Y
    l0 = L0

    # 3 dummy sensors (no cylinders — just empty channel)
    for y_off in [2.0, 0.0, -2.0]:
        ff.add_sensor((40.0 * l0, CENTER_Y + y_off * l0, 0.0), l0 / 4.0)

    n_obj = 3
    stabilize = int(4 * 1280 / cfg["u0"])
    print(f"Stabilising empty channel ({stabilize} steps)...")
    ff.run(stabilize, np.zeros(n_obj, dtype=np.float32))

    n_steps = 100  # steady-state, 100 is more than enough
    si = cfg["sample_interval"]
    sens_list = []
    ux_list, uy_list = [], []

    for step in range(n_steps):
        ff.run(si, np.zeros(n_obj, dtype=np.float32))
        obs = ff.obs.copy()
        sens_list.append(obs[0:6])
        ux, uy = get_velocity_field(ff, u0=cfg["u0"])
        ux_list.append(ux)
        uy_list.append(uy)

    # Save
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.stack(ux_list), uy=np.stack(uy_list))
    np.savez(os.path.join(out_dir, "sensors.npz"),
             sensors=np.array(sens_list, dtype=np.float32))

    # Validate: should be nearly steady
    sens_arr = np.array(sens_list, dtype=np.float32)
    print(f"  Sensor mean: {np.mean(sens_arr, axis=0)}")
    print(f"  Sensor std:  {np.std(sens_arr, axis=0)}")
    total_std = float(np.sqrt(np.mean(sens_arr ** 2)))
    print(f"  Total std (should be near 0 for steady channel): {total_std:.6f}")

    # Vorticity
    omega = vorticity_from_ddf(ff, u0=cfg["u0"])
    save_vorticity_png(os.path.join(out_dir, "vorticity.png"),
                       omega, title="Empty channel Re=100")

    # Meta
    meta = dict(cfg, n_steps=n_steps)
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)
    print(f"Empty channel data saved to {out_dir}")

    del ff


if __name__ == "__main__":
    t0 = time.time()
    collect()
    print(f"Time: {time.time() - t0:.1f}s")
