# CelerisLab/tests/validation/test_sensor_accuracy.py
"""Sensor accuracy validation: compare sensor readings to direct flow field averages.

This script validates that the GPU sensor kernel accumulation matches a
CPU-side manual average of the macroscopic field over the same cell footprint.

Usage::

    conda run -n pycuda_3_10 python tests/validation/test_sensor_accuracy.py
"""

from __future__ import annotations

import json
import os
import sys
import tempfile
from pathlib import Path

import numpy as np

_REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(_REPO / "src"))

from CelerisLab import Simulation


def test_sensor_accuracy() -> dict:
    """Run sensor accuracy validation with multiple sensor positions."""
    cfg = json.loads(
        (Path(_REPO) / "src" / "CelerisLab" / "configs" / "config_lbm.json").read_text()
    )
    cfg["grid"]["nx"] = 256
    cfg["grid"]["ny"] = 128
    cfg["grid"]["nz"] = 1
    cfg["physics"]["viscosity"] = 0.009
    cfg["physics"]["velocity"] = 0.03
    cfg["method"]["collision"] = "MRT"
    cfg["method"]["inlet"]["scheme"] = "regularized"
    cfg["method"]["inlet"]["profile"] = "uniform"
    cfg["method"]["y_wall_bc"] = "free_slip"

    tmpd = tempfile.mkdtemp(prefix="sensor_test_")
    lbm_path = os.path.join(tmpd, "config_lbm.json")
    with open(lbm_path, "w") as f:
        json.dump(cfg, f)

    sim = Simulation(lbm_config_path=lbm_path)
    sim.add_body("circle", center=(80, 64), radius=15)

    positions = [(120, 50), (120, 64), (120, 78), (150, 64)]
    sensor_ids = []
    for cx, cy in positions:
        sid = sim.add_body("sensor", center=(cx, cy), radius=10)
        sensor_ids.append(sid)

    sim.initialize()
    print(f"Initialized: nx={cfg['grid']['nx']} ny={cfg['grid']['ny']} "
          f"n_curved={sim.field.n_curved} n_sensor={sim.field.n_sensor}")

    # Step to develop wake
    for _ in range(50):
        sim.run(20)

    # Get macroscopic field after one more step (with sensor accumulation)
    import pycuda.driver as cuda
    stream = cuda.Stream()
    sim.bodies.zero_sensor_segment_async(stream)
    sim.stepper.step(1, action_gpu=sim.bodies.action_gpu,
                     obs_gpu=sim.bodies.obs_gpu, stream=stream)
    stream.synchronize()

    macro = sim.get_macroscopic()
    ux = macro["ux"]
    uy = macro["uy"]

    results = {}
    all_pass = True
    for sid in sensor_ids:
        cells_arr, _ = sim.bodies.get(sid).get_sensor_list(
            sim.lbm_cfg.nx, sim.lbm_cfg.ny
        )
        cell_idx = np.asarray(cells_arr, dtype=np.int64)
        ux_rav = ux.ravel().astype(np.float64)
        uy_rav = uy.ravel().astype(np.float64)

        sensor_ux_mean = float(np.mean(ux_rav[cell_idx]))
        sensor_uy_mean = float(np.mean(uy_rav[cell_idx]))

        sensor_reading = sim.read_sensor(sid)
        sensor_reading_x = float(sensor_reading[0])
        sensor_reading_y = float(sensor_reading[1])

        diff_ux = abs(sensor_reading_x - sensor_ux_mean)
        diff_uy = abs(sensor_reading_y - sensor_uy_mean)
        passed = diff_ux < 1e-4 and diff_uy < 1e-4
        if not passed:
            all_pass = False

        results[f"sensor_{sid}_pos{positions[i]}"] = {
            "sensor_reading": [sensor_reading_x, sensor_reading_y],
            "manual_average": [sensor_ux_mean, sensor_uy_mean],
            "diff": [float(diff_ux), float(diff_uy)],
            "n_cells": int(len(cells_arr)),
            "pass": bool(passed),
        }
        status = "PASS" if passed else "FAIL"
        print(
            f"  Sensor {sid} @ {positions[sid]}: "
            f"reading=({sensor_reading_x:.8f},{sensor_reading_y:.8f}) "
            f"manual=({sensor_ux_mean:.8f},{sensor_uy_mean:.8f}) "
            f"diff=({diff_ux:.2e},{diff_uy:.2e}) "
            f"cells={len(cells_arr)} [{status}]"
        )

    sim.close()
    summary = {"all_pass": bool(all_pass), "results": results}
    print(f"\nSensor accuracy: {'ALL PASS' if all_pass else 'SOME FAILED'}")
    return summary


if __name__ == "__main__":
    result = test_sensor_accuracy()
    sys.exit(0 if result["all_pass"] else 1)