"""Parity + memory gate: fused Metal sparse conv vs the stock chunked path.
Small case first (exact check vs naive reference), then a decoder-scale
case for peak-memory + timing comparison.
"""
import sys
import time
from pathlib import Path

import mlx.core as mx
import numpy as np

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

from mlx_backend.sparse_conv import MlxSparseConv3d, _build_neighbor_map
from mlx_backend.sparse_conv_metal import sparse_conv_metal
from mlx_backend.sparse_tensor import MlxSparseTensor

mx.random.seed(0)


def make_case(n_vox, ci, co, res):
    # random unique voxel coords in a res^3 grid, batch 0
    rng = np.random.default_rng(0)
    flat = rng.choice(res ** 3, size=n_vox, replace=False)
    x, rem = flat // (res * res), flat % (res * res)
    y, z = rem // res, rem % res
    coords = np.stack([np.zeros_like(x), x, y, z], 1).astype(np.int32)
    st = MlxSparseTensor(
        feats=mx.random.normal((n_vox, ci)).astype(mx.float16),
        coords=mx.array(coords))
    conv = MlxSparseConv3d(ci, co, 3)
    conv.weight = mx.random.normal((co, 3, 3, 3, ci)).astype(mx.float16) * 0.05
    conv.bias = mx.random.normal((co,)).astype(mx.float16) * 0.1
    return st, conv


def run_metal(st, conv):
    nmap = _build_neighbor_map(st.coords, st.shape[0], st.spatial_shape,
                               conv.kernel_size, conv.dilation)
    mx.eval(nmap)
    K = nmap.shape[0]
    Co, Ci = conv.out_channels, conv.in_channels
    w = conv.weight.reshape(Co, K, Ci).transpose(1, 2, 0)  # (K, Ci, Co)
    feats_padded = mx.concatenate(
        [st.feats, mx.zeros((1, Ci), dtype=st.feats.dtype)], axis=0)
    return sparse_conv_metal(feats_padded, nmap, w, conv.bias)


# ---- small exact-parity case ----
st, conv = make_case(5000, 64, 96, 64)
ref = conv(st).feats
got = run_metal(st, conv)
mx.eval(ref, got)
d = mx.max(mx.abs(ref.astype(mx.float32) - got.astype(mx.float32))).item()
rel = d / max(1e-6, mx.max(mx.abs(ref.astype(mx.float32))).item())
print(f"small: max|diff|={d:.3e} rel={rel:.3e}", "PASS" if rel < 2e-2 else "FAIL")
ok_small = rel < 2e-2

# ---- decoder-scale case: memory + time ----
st2, conv2 = make_case(400_000, 512, 512, 512)
mx.eval(st2.feats, conv2.weight)

mx.reset_peak_memory()
t0 = time.perf_counter()
ref2 = conv2(st2).feats
mx.eval(ref2)
t_stock = time.perf_counter() - t0
m_stock = mx.get_peak_memory() / 2**30

st2._spatial_cache = {}  # drop cached neighbor map so both build it
mx.reset_peak_memory()
t0 = time.perf_counter()
got2 = run_metal(st2, conv2)
mx.eval(got2)
t_metal = time.perf_counter() - t0
m_metal = mx.get_peak_memory() / 2**30

d2 = mx.max(mx.abs(ref2.astype(mx.float32) - got2.astype(mx.float32))).item()
r2 = d2 / max(1e-6, mx.max(mx.abs(ref2.astype(mx.float32))).item())
print(f"large: stock {t_stock:.2f}s peak {m_stock:.2f}GB | "
      f"metal {t_metal:.2f}s peak {m_metal:.2f}GB | rel diff {r2:.3e}")
ok_large = r2 < 2e-2
print("SPCONV_GATE:", "PASS" if (ok_small and ok_large) else "FAIL")
sys.exit(0 if (ok_small and ok_large) else 1)