"""Reference implementation of the sar3d volume from Multi_antenna_3D_clean_v1.ipynb.

Executable specification for the C++ port: recomputes a single-frame
(aperture radius 0) 3D SAR volume from a preprocessed snapshot exactly the
way the notebook does, then compares it against a reference *_volume.npz
exported by the notebook.

Usage:
    python sar3d_reference_check.py <dataset_root> <reference_volume.npz>

<dataset_root> is an extracted replay dataset (contains config_profile.json
and preprocessed/<snapshot>/i*_o*_{s21,freq}.npy + meta.json). The snapshot
name and all physics parameters are taken from the npz meta.

Exits non-zero when the recomputed volume differs from the reference by more
than the tolerance. The only expected difference is the Gaussian smoothing
(manual separable filter vs scipy) at the 1e-6 level.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

TOLERANCE = 2e-6

C_LIGHT = 3e8

# Notebook constants not recorded in the npz meta (cell 0 / cell 4).
BP_OVERSAMPLE = 8
BP_WINDOW = True
PAIR_NORM_PERCENTILE = 50.0
PAIR_NORM_EPS = 1e-15
COMP_RANGE_POWER = 0.3
COMP_RANGE_WEIGHT_MAX = 5.0
COMP_REF_DEPTH = 5.0
BP_SMOOTH_SIGMA = 1.5
SAR_Y_RANGE_M = (-0.5, 2.0)
SAR_NORM_EDGE_MARGIN_M = 0.3
X_MARGIN_M = 2.0
MOTION_PHASE_SIGN = -1.0
DIRECTION_SIGN = +1.0


def gaussian_kernel1d(sigma: float, radius: int) -> np.ndarray:
    x = np.arange(-radius, radius + 1, dtype=float)
    phi = np.exp(-0.5 / (sigma * sigma) * x * x)
    return phi / phi.sum()


def correlate1d_reflect(arr: np.ndarray, weights: np.ndarray, axis: int) -> np.ndarray:
    r = (len(weights) - 1) // 2
    a = np.moveaxis(arr, axis, -1)
    pad = np.concatenate([a[..., r - 1::-1], a, a[..., :-r - 1:-1]], axis=-1)
    out = np.zeros_like(a)
    for k, w in enumerate(weights):
        out += w * pad[..., k:k + a.shape[-1]]
    return np.moveaxis(out, -1, axis)


def gaussian_filter(arr: np.ndarray, sigma: tuple[float, ...], truncate: float = 4.0) -> np.ndarray:
    """scipy.ndimage.gaussian_filter equivalent (mode='reflect')."""
    out = np.asarray(arr, dtype=float).copy()
    for axis, s in enumerate(sigma):
        if s <= 0:
            continue
        radius = int(truncate * float(s) + 0.5)
        out = correlate1d_reflect(out, gaussian_kernel1d(float(s), radius), axis)
    return out


def load_geometry(profile_path: Path) -> dict:
    profile = json.loads(profile_path.read_text())
    gpr = profile["gpr"]
    tx_geom = {int(e["output_pos"]): (e["x_m"], e.get("y_m", 0.0), e.get("z_m", 0.0))
               for e in gpr["tx_geometry"]}
    rx_geom = {int(e["input_pos"]): (e["x_m"], e.get("y_m", 0.0), e.get("z_m", 0.0))
               for e in gpr["rx_geometry"]}
    combos = profile["run"]["combos"]
    pair_order = [(int(c["output"]), int(c["input"])) for c in combos]
    tx_order_full: list[int] = []
    for tx, _ in pair_order:
        if tx not in tx_order_full:
            tx_order_full.append(tx)
    output_idx = sorted(tx_geom)
    input_idx = sorted(rx_geom)
    tx_xyz = np.array([tx_geom[i] for i in output_idx], dtype=float)
    rx_xyz = np.array([rx_geom[i] for i in input_idx], dtype=float)
    return {
        "output_idx": output_idx, "input_idx": input_idx,
        "tx_xyz": tx_xyz, "rx_xyz": rx_xyz,
        "pair_order": pair_order, "tx_order_full": tx_order_full,
    }


def snapshot_timing(snapshot_dir: Path) -> dict | None:
    """Per-Tx capture windows from the snapshot's meta.json trace labels."""
    try:
        meta = json.loads((snapshot_dir / "meta.json").read_text())
    except OSError:
        return None
    traces = [t for t in meta.get("traces", [])
              if "capture_start_ns" in t and "capture_end_ns" in t and "output" in t]
    if not traces:
        return None
    first = min(int(t["capture_start_ns"]) for t in traces)
    last = max(int(t["capture_end_ns"]) for t in traces)
    by_tx: dict[int, tuple[int, int]] = {}
    for t in traces:
        by_tx.setdefault(int(t["output"]), (int(t["capture_start_ns"]), int(t["capture_end_ns"])))
    return {
        "acq_s": (last - first) * 1e-9,
        "center_abs_s": 0.5 * (first + last) * 1e-9,
        "tx_window_s": {tx: ((a - first) * 1e-9, (b - first) * 1e-9) for tx, (a, b) in by_tx.items()},
    }


def pair_timestamps(geometry: dict, timing: dict, speed_m_s: float, look_angle_deg: float,
                    wave_speed: float) -> dict:
    """Segment/slot assignment and coarse pair dtau (notebook compute_pair_timestamps_for_snapshot)."""
    windows = timing["tx_window_s"]
    hw_order = [tx for tx in geometry["tx_order_full"] if tx in windows]
    hw_order += [tx for tx in sorted(windows) if tx not in hw_order]
    groups: dict[tuple[float, float], list[int]] = {}
    for tx in hw_order:
        groups.setdefault(windows[tx], []).append(tx)
    seg_of = {}
    for seg_idx, win in enumerate(sorted(groups)):
        for slot, tx in enumerate(groups[win]):
            seg_of[tx] = (seg_idx, slot, len(groups[win]))
    cos_theta = np.cos(np.radians(look_angle_deg))
    t_ref = 0.5 * timing["acq_s"]
    tx_map = {p: k for k, p in enumerate(geometry["output_idx"])}
    rx_map = {p: k for k, p in enumerate(geometry["input_idx"])}
    ts = {}
    for tx_phys, rx_phys in geometry["pair_order"]:
        _, slot, n_slots = seg_of[tx_phys]
        t0, t1 = windows[tx_phys]
        t_center = 0.5 * (t0 + t1)
        dz_motion = DIRECTION_SIGN * speed_m_s * (t_center - t_ref) * cos_theta
        ts[(tx_map[tx_phys], rx_map[rx_phys])] = {
            "slot": slot, "n_slots": n_slots,
            "segment_start_s": t0, "segment_time_s": t1 - t0,
            "t_center_s": t_center,
            "dtau_motion_s": 2.0 * dz_motion / wave_speed,
        }
    return ts


def timing_abs_center(snapshot_dir: Path) -> float:
    timing = snapshot_timing(snapshot_dir)
    if timing is None:
        raise RuntimeError(f"meta.json trace labels required in {snapshot_dir}")
    return timing["center_abs_s"]


def load_mimo(snapshot_dir: Path, geometry: dict) -> tuple[dict, dict]:
    s21_data, freq_data = {}, {}
    for f in snapshot_dir.glob("i*_o*_s21.npy"):
        parts = f.stem.split("_")
        rx_phys = int(parts[0][1:])
        tx_phys = int(parts[1][1:])
        if tx_phys not in geometry["output_idx"] or rx_phys not in geometry["input_idx"]:
            continue
        key = (geometry["output_idx"].index(tx_phys), geometry["input_idx"].index(rx_phys))
        s21_data[key] = np.load(f)
        freq_data[key] = np.asarray(np.load(f.parent / f"i{rx_phys}_o{tx_phys}_freq.npy"), dtype=float)
    return s21_data, freq_data


def apply_phase_correction(s21: np.ndarray, freq_full: np.ndarray, info: dict,
                           f_start: float, f_stop: float, speed_m_s: float,
                           look_angle_deg: float, wave_speed: float) -> np.ndarray:
    """Coarse pair + intra-sweep motion correction as one phase ramp before the IFFT."""
    idx = np.flatnonzero((freq_full >= f_start) & (freq_full <= f_stop))
    n_full = len(freq_full)
    dt_event = info["segment_time_s"] / (info["n_slots"] * n_full)
    t_abs = info["segment_start_s"] + (info["n_slots"] * idx + info["slot"]) * dt_event
    dt_intra = t_abs - info["t_center_s"]
    dtau_intra = 2.0 * DIRECTION_SIGN * speed_m_s * np.cos(np.radians(look_angle_deg)) * dt_intra / wave_speed
    total_dtau = dtau_intra + info["dtau_motion_s"]
    phi = 2.0 * np.pi * freq_full[idx] * total_dtau
    out = np.array(s21, dtype=np.complex128, copy=True)
    out[idx] *= np.exp(1j * MOTION_PHASE_SIGN * phi)
    return out


def compute_ascan(s21: np.ndarray, freq: np.ndarray, f_start: float, f_stop: float):
    """Oversampled complex A-scan with the notebook's frequency-shifted IFFT."""
    mask = (freq >= f_start) & (freq <= f_stop)
    freq_cut = freq[mask]
    s21_cut = s21[mask]
    df = float(np.median(np.diff(freq_cut)))
    n = len(freq_cut)
    k0 = int(round(freq_cut[0] / df))
    n_fft = (1 << int(np.ceil(np.log2(2 * (k0 + n - 1))))) * BP_OVERSAMPLE
    dt = 1.0 / (n_fft * df)
    s = s21_cut * np.hanning(n) if BP_WINDOW else s21_cut.copy()
    spectrum = np.zeros(n_fft, dtype=np.complex128)
    spectrum[k0:k0 + n] = s
    return np.arange(n_fft, dtype=float) * dt, np.fft.ifft(spectrum)


def interp_complex(tau: np.ndarray, t_axis: np.ndarray, h_axis: np.ndarray) -> np.ndarray:
    n = t_axis.size
    dt = (float(t_axis[-1]) - float(t_axis[0])) / (n - 1)
    pos = (tau - float(t_axis[0])) / dt
    inside = (pos >= 0.0) & (pos <= n - 1)
    k = np.clip(np.floor(pos), 0, n - 2).astype(np.intp)
    frac = pos - k
    out = h_axis[k] + frac * (h_axis[k + 1] - h_axis[k])
    return np.where(inside, out, 0.0 + 0.0j)


def axis_linspace(lo: float, hi: float, step: float) -> np.ndarray:
    return np.linspace(lo, hi, max(int(round((hi - lo) / step)) + 1, 2))


def compute_volume(dataset_root: Path, meta: dict) -> dict:
    """Single-frame volume via the notebook's ANIMATION path.

    The reference npz files were exported from animation frames: the base grid
    is shifted by the radar displacement between the frame and the scene-origin
    snapshot (meta reference_snapshot), antennas get the same offset, and the
    depth window follows the radar (depth_shift). Physically it all cancels for
    an aperture of radius 0, but replicating the exact float arithmetic is what
    makes the comparison bit-tight at the depth-mask knife edge.
    """
    geometry = load_geometry(dataset_root / "config_profile.json")
    snapshot_dir = dataset_root / "preprocessed" / meta["snapshot"]

    eps_r = 1.0
    wave_speed = C_LIGHT / np.sqrt(eps_r)
    f_start = float(meta["f_start_hz"])
    f_stop = float(meta["f_stop_hz"])
    min_depth = float(meta["min_depth_m"])
    max_depth = float(meta["max_depth_m"])
    speed = float(meta["speed_m_s"])
    look_angle = float(meta["look_angle_deg"])
    ground_depth = float(meta["ground_depth_m"])
    voxel = [float(a) for a in meta["voxel_m"]]

    x_tx, y_tx, z_tx = geometry["tx_xyz"].T
    x_rx, y_rx, z_rx = geometry["rx_xyz"].T
    n_tx, n_rx = len(x_tx), len(x_rx)

    timing = snapshot_timing(snapshot_dir)
    if timing is None:
        raise RuntimeError(f"meta.json trace labels required in {snapshot_dir}")
    ts = pair_timestamps(geometry, timing, speed, look_angle, wave_speed)

    s21_data, freq_data = load_mimo(snapshot_dir, geometry)
    ascan_h, ascan_t = {}, {}
    for key, s21 in s21_data.items():
        corrected = apply_phase_correction(
            s21, freq_data[key], ts[key], f_start, f_stop, speed, look_angle, wave_speed)
        t_axis, h = compute_ascan(corrected, freq_data[key], f_start, f_stop)
        z_axis = t_axis * wave_speed / 2
        m = (z_axis >= min_depth) & (z_axis <= max_depth)
        scale = float(np.percentile(np.abs(h)[m], PAIR_NORM_PERCENTILE))
        if not np.isfinite(scale) or scale <= PAIR_NORM_EPS:
            scale = 1.0
        ascan_h[key] = h / (scale + PAIR_NORM_EPS)
        ascan_t[key] = t_axis

    pitch = np.radians(look_angle)
    cos_p, sin_p = np.cos(pitch), np.sin(pitch)

    def scene_to_radar_yz(y_scene, z_scene):
        y0 = y_scene - ground_depth
        return y0 * cos_p + z_scene * sin_p, -y0 * sin_p + z_scene * cos_p

    x_ant = np.concatenate([x_tx, x_rx])
    y_lo, y_hi = SAR_Y_RANGE_M
    z_span = [(d - float(scene_to_radar_yz(y, 0.0)[1])) / cos_p
              for y in (y_lo, y_hi) for d in (min_depth, max_depth)]
    gx = axis_linspace(float(x_ant.min() - X_MARGIN_M), float(x_ant.max() + X_MARGIN_M), voxel[0])
    gy = axis_linspace(y_lo, y_hi, voxel[1])
    gz = axis_linspace(min(z_span), max(z_span), voxel[2])

    # Displacement of this frame relative to the scene-origin snapshot
    # (notebook sar_frame_offsets + anim_frame_displacement).
    off_dy = off_dz = 0.0
    disp_y = disp_z = 0.0
    reference_name = meta.get("reference_snapshot", meta["snapshot"])
    if reference_name != meta["snapshot"]:
        ref_timing = snapshot_timing(dataset_root / "preprocessed" / reference_name)
        if ref_timing is None:
            raise RuntimeError(f"meta.json trace labels required in {reference_name}")
        dt = timing_abs_center(snapshot_dir) - ref_timing["center_abs_s"]
        horizontal = DIRECTION_SIGN * speed * dt
        off_dy = float(horizontal * np.sin(pitch))
        off_dz = float(horizontal * np.cos(pitch))
        # Scene displacement = difference of two transformed points; the ground
        # lift cancels, leaving essentially (0, ~0, horizontal).
        disp_y = (off_dy * cos_p - off_dz * sin_p + ground_depth) - ground_depth
        disp_z = off_dy * sin_p + off_dz * cos_p
    gy_s = gy + disp_y
    gz_s = gz + disp_z
    depth_shift = off_dz
    grid_xx, grid_zz = np.meshgrid(gx, gz_s)

    ground_plane_y = (-ground_depth + COMP_REF_DEPTH * sin_p) / cos_p

    def pair_geo_ref(i, j):
        xc = (x_tx[i] + x_rx[j]) / 2.0
        r_tx = np.sqrt((xc - x_tx[i]) ** 2 + (ground_plane_y - y_tx[i]) ** 2 + (COMP_REF_DEPTH - z_tx[i]) ** 2)
        r_rx = np.sqrt((xc - x_rx[j]) ** 2 + (ground_plane_y - y_rx[j]) ** 2 + (COMP_REF_DEPTH - z_rx[j]) ** 2)
        return 1.0 / (r_tx * r_rx + 1e-12) + 1e-30

    geo_ref = {(i, j): pair_geo_ref(i, j) for i in range(n_tx) for j in range(n_rx)}
    dx2_tx = [(grid_xx - x_tx[i]) ** 2 for i in range(n_tx)]
    dx2_rx = [(grid_xx - x_rx[j]) ** 2 for j in range(n_rx)]

    acc_vol = np.zeros((gy.size, gz.size, gx.size), dtype=np.complex128)
    cnt_vol = np.zeros(acc_vol.shape, dtype=float)
    for iy, y_plane in enumerate(gy_s):
        y_rad, z_rad = scene_to_radar_yz(y_plane, grid_zz)
        depth_ok = (z_rad - depth_shift >= min_depth) & (z_rad - depth_shift <= max_depth)
        r_tx_all = [np.sqrt(dx2_tx[i] + (y_rad - (y_tx[i] + off_dy)) ** 2 + (z_rad - (z_tx[i] + off_dz)) ** 2)
                    for i in range(n_tx)]
        r_rx_all = [np.sqrt(dx2_rx[j] + (y_rad - (y_rx[j] + off_dy)) ** 2 + (z_rad - (z_rx[j] + off_dz)) ** 2)
                    for j in range(n_rx)]
        acc = np.zeros_like(grid_xx, dtype=np.complex128)
        cnt = np.zeros_like(grid_xx, dtype=float)
        for i in range(n_tx):
            for j in range(n_rx):
                key = (i, j)
                if key not in ascan_h:
                    continue
                r_tx, r_rx = r_tx_all[i], r_rx_all[j]
                tau = (r_tx + r_rx) / wave_speed
                t_axis = ascan_t[key]
                valid = depth_ok & (tau >= t_axis[0]) & (tau <= t_axis[-1])
                h = interp_complex(tau, t_axis, ascan_h[key])
                ratio = (1.0 / (r_tx * r_rx + 1e-12) + 1e-30) / geo_ref[key]
                weight = np.clip(1.0 / (ratio ** COMP_RANGE_POWER + 1e-12), 0.0, COMP_RANGE_WEIGHT_MAX)
                acc += np.where(valid, h * weight, 0.0 + 0.0j)
                cnt += valid
        acc_vol[iy] = acc
        cnt_vol[iy] = cnt

    volume = np.abs(acc_vol / (cnt_vol + 1e-12))
    z_rad_planes = np.asarray([scene_to_radar_yz(yp, grid_zz)[1] - depth_shift for yp in gy_s])
    volume = np.where((z_rad_planes >= min_depth) & (z_rad_planes <= max_depth), volume, 0.0)
    ref_planes = (gy_s >= gy_s[0] + SAR_NORM_EDGE_MARGIN_M) & (gy_s <= gy_s[-1] - SAR_NORM_EDGE_MARGIN_M)
    level_raw = volume[ref_planes].max() + 1e-12
    volume = gaussian_filter(volume / level_raw, sigma=(0.0, BP_SMOOTH_SIGMA, BP_SMOOTH_SIGMA))
    level_smooth = volume[ref_planes].max() + 1e-12
    volume = np.clip(volume / level_smooth, 0.0, None).astype(np.float32)
    return {"volume": volume, "level": float(level_raw * level_smooth), "x": gx, "y": gy, "z": gz}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("dataset_root", type=Path,
                        help="extracted dataset dir (config_profile.json + preprocessed/)")
    parser.add_argument("reference_npz", type=Path, help="reference *_volume.npz from the notebook")
    args = parser.parse_args()

    ref = np.load(args.reference_npz)
    meta = json.loads(str(ref["meta"]))
    print(f"snapshot: {meta['snapshot']}")
    result = compute_volume(args.dataset_root, meta)

    volume, ref_volume = result["volume"], ref["volume"]
    axes_ok = all(np.allclose(result[k], ref[f"{k}_m"]) for k in ("x", "y", "z"))
    level_rel = abs(result["level"] - meta["level"]) / abs(meta["level"])
    print(f"axes match: {axes_ok}; level rel diff: {level_rel:.3e} "
          f"(mine {result['level']:.9f}, ref {meta['level']:.9f})")
    if volume.shape != ref_volume.shape:
        print(f"FAIL: shape mismatch {volume.shape} vs {ref_volume.shape}")
        return 1
    diff = np.abs(volume.astype(np.float64) - ref_volume.astype(np.float64))
    argmax = np.unravel_index(int(np.argmax(volume)), volume.shape)
    ref_argmax = np.unravel_index(int(np.argmax(ref_volume)), ref_volume.shape)
    print(f"max |diff| = {diff.max():.3e}, mean = {diff.mean():.3e}")
    print(f"argmax: mine {argmax}, ref {ref_argmax}")
    ok = axes_ok and diff.max() <= TOLERANCE and argmax == ref_argmax and level_rel <= TOLERANCE
    print("OK" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())