"""Compare a sar3d volume dumped by the C++ processor with a notebook npz export.

Usage:
    python sar3d_volume_compare.py <dump_base> <reference_volume.npz>

<dump_base> is the path prefix written by `sar3d_parity --dump-dir` (or by the
processor with `sar3d_volume_dump_dir` set): <dump_base>.f32 holds the raw
float32 volume [iy][iz][ix], <dump_base>.json its axes and level.

The C++ path accumulates the backprojection in float32 (CUDA kernel), so the
comparison is tolerance-based rather than bitwise; geometry is computed in
double, keeping the expected difference at the 1e-5 level. Voxels within the
smoothing radius of the depth-mask boundary are reported separately: there the
float32 knife-edge can flip whole boundary voxels (the same effect exists
between two float64 runs with shifted grids).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

MAX_DIFF_TOLERANCE = 1e-4
LEVEL_REL_TOLERANCE = 1e-4
# Gaussian smoothing radius (sigma 1.5 -> radius 6) around the depth-mask edge.
EDGE_EXCLUSION_VOXELS = 7


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("dump_base", type=Path, help="path prefix of the .f32/.json dump")
    parser.add_argument("reference_npz", type=Path, help="reference *_volume.npz from the notebook")
    args = parser.parse_args()

    header = json.loads(Path(str(args.dump_base) + ".json").read_text())
    shape = (header["ny"], header["nz"], header["nx"])
    volume = np.fromfile(str(args.dump_base) + ".f32", dtype=np.float32).reshape(shape)

    ref = np.load(args.reference_npz)
    ref_volume = ref["volume"]
    meta = json.loads(str(ref["meta"]))
    print(f"snapshot: {meta['snapshot']}")
    if volume.shape != ref_volume.shape:
        print(f"FAIL: shape mismatch {volume.shape} vs {ref_volume.shape}")
        return 1

    axes_ok = all(
        np.allclose(np.asarray(header[key]), ref[key])
        for key in ("x_m", "y_m", "z_m")
    )
    level_rel = abs(header["level"] - meta["level"]) / abs(meta["level"])
    print(f"axes match: {axes_ok}; level rel diff: {level_rel:.3e} "
          f"(dump {header['level']:.9f}, ref {meta['level']:.9f})")

    diff = np.abs(volume.astype(np.float64) - ref_volume.astype(np.float64))

    # Depth-mask boundary band: voxels whose (iy, iz) sit within the smoothing
    # radius of a masked/unmasked transition. Reconstructed from the reference
    # parameters exactly like the processor computes the mask.
    pitch = np.radians(float(meta.get("ground_pitch_deg") or meta["look_angle_deg"]))
    y0 = np.asarray(ref["y_m"]) - float(meta["ground_depth_m"])
    z_rad = -y0[:, None] * np.sin(pitch) + np.asarray(ref["z_m"])[None, :] * np.cos(pitch)
    mask = (z_rad >= float(meta["min_depth_m"])) & (z_rad <= float(meta["max_depth_m"]))
    edge = np.zeros_like(mask)
    r = EDGE_EXCLUSION_VOXELS
    for shift in range(-r, r + 1):
        edge |= np.roll(mask, shift, axis=1) != mask
    interior = ~edge

    argmax = np.unravel_index(int(np.argmax(volume)), volume.shape)
    ref_argmax = np.unravel_index(int(np.argmax(ref_volume)), ref_volume.shape)
    interior_diff = diff[interior[:, :, None] & np.ones(volume.shape, dtype=bool)]
    print(f"max |diff| (all voxels)          = {diff.max():.3e}")
    print(f"max |diff| (off depth-mask edge) = {interior_diff.max():.3e}, "
          f"mean = {interior_diff.mean():.3e}")
    print(f"argmax: dump {argmax}, ref {ref_argmax}")

    ok = (axes_ok and level_rel <= LEVEL_REL_TOLERANCE and argmax == ref_argmax
          and interior_diff.max() <= MAX_DIFF_TOLERANCE)
    print("OK" if ok else "FAIL")
    return 0 if ok else 1


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