"""Behavioural tests for the two sparse decoders.

These are the only Pixal3D models using sparse convolution, so upstream cannot run here
(spconv has no Metal build) and there is NO numerical oracle — unlike the flow models and
ss_dec, which are all verified at correlation 1.0 against upstream on CPU torch.

So these check structure and behaviour rather than values: that weights map completely,
that the occupied set grows SELECTIVELY rather than x8, that scale tracks the four
upsamples, and that the vertex head lands inside the band its sigmoid+margin allows.
Weaker than a correlation, and stated as such.
"""

import sys
from pathlib import Path

import mlx.core as mx
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from pixal3d_mlx.decoders import load  # noqa: E402
from trellis_sparse_mlx import SparseTensor  # noqa: E402

CK = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
W = Path(__file__).resolve().parents[1] / "weights"


def _latent(n=12, ch=32, res=4, seed=0):
    rng = np.random.default_rng(seed)
    co = np.concatenate(
        [np.zeros((n, 1), dtype=np.int32), rng.integers(0, res, (n, 3)).astype(np.int32)], 1
    )
    co = np.unique(co, axis=0)
    return SparseTensor(
        mx.array(rng.standard_normal((len(co), ch)).astype(np.float32)), mx.array(co)
    )


def _have(stem):
    return (W / f"{stem}.safetensors").exists() and (CK / f"{stem}.json").exists()


def test_shape_dec():
    stem = "shape_dec_next_dc_f16c32_fp16"
    if not _have(stem):
        print("      (skipped: weights absent)")
        return 0.0
    m, rep = load(W / f"{stem}.safetensors", CK / f"{stem}.json")
    assert not rep["missing"] and not rep["unmapped"] and not rep["mismatched"], rep
    x = _latent()
    n_in = x.coords.shape[0]
    out, subs = m(x, return_subs=True)
    f = np.asarray(out.feats)

    assert f.shape[1] == 7, f"FlexiDualGrid emits 7 channels, got {f.shape[1]}"
    assert np.isfinite(f).all(), "non-finite output"
    # four x2 upsamples -> scale 1/16
    assert abs(out._scale[0] - 1 / 16) < 1e-9, out._scale
    # growth must be selective, not the full x8 per stage
    assert out.coords.shape[0] > n_in, "no growth at all"
    assert out.coords.shape[0] < n_in * 8**4, "grew by the full x8 every stage"
    assert len(subs) == 4, f"expected 4 subdivision masks, got {len(subs)}"

    v = np.asarray(m.decode_vertices(out))
    m_ = m.voxel_margin
    assert v.min() >= -m_ - 1e-4 and v.max() <= 1 + m_ + 1e-4, (v.min(), v.max())
    return 0.0


def test_tex_dec_follows_guide_masks():
    """tex_dec has pred_subdiv=False, so its structure must match the masks it is given."""
    s_stem, t_stem = "shape_dec_next_dc_f16c32_fp16", "tex_dec_next_dc_f16c32_fp16"
    if not (_have(s_stem) and _have(t_stem)):
        print("      (skipped: weights absent)")
        return 0.0
    sm, _ = load(W / f"{s_stem}.safetensors", CK / f"{s_stem}.json")
    tm, rep = load(W / f"{t_stem}.safetensors", CK / f"{t_stem}.json")
    assert not rep["missing"] and not rep["unmapped"], rep

    x = _latent()
    shape_out, subs = sm(x, return_subs=True)
    tex_out = tm(_latent(), guide_subs=subs)
    f = np.asarray(tex_out.feats)

    assert f.shape[1] == 6, f"tex_dec emits 6 channels, got {f.shape[1]}"
    assert np.isfinite(f).all(), "non-finite output"
    assert tex_out.coords.shape[0] == shape_out.coords.shape[0], (
        "guided decoder produced a different voxel count than the masks describe"
    )
    return 0.0


if __name__ == "__main__":
    tests = [
        ("shape_dec structure", test_shape_dec),
        ("tex_dec follows guides", test_tex_dec_follows_guide_masks),
    ]
    failed = 0
    for name, fn in tests:
        try:
            fn()
            print(f"  PASS  {name}")
        except AssertionError as e:
            print(f"  FAIL  {name}: {e}")
            failed += 1
        except Exception as e:  # noqa: BLE001
            print(f"  ERROR {name}: {type(e).__name__}: {e}")
            failed += 1
    print(f"\n{len(tests)-failed}/{len(tests)} passed")
    sys.exit(1 if failed else 0)