#!/usr/bin/env python3
"""test_betti_tracker.py — unit tests for betti_tracker.py

Run:
  source /tmp/betti_env/bin/activate && python3 test_betti_tracker.py
"""

from __future__ import annotations
import json, sys, tempfile, warnings
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from betti_tracker import (
    velocity_gradient,
    scar_density,
    threshold_scar_support,
    compute_betti_numbers,
    run_pipeline,
)

warnings.filterwarnings("ignore")


def _synthetic_field(N: int, T: int, void: bool = False):
    u = np.ones((T, N, N, N, 3), dtype=np.float64) * 0.5
    if void:
        cy = cx = cz = N // 2
        r = N // 6
        Y, X, Z = np.ogrid[:N, :N, :N]
        sphere = (X - cx) ** 2 + (Y - cy) ** 2 + (Z - cz) ** 2 <= r ** 2
        u[:, sphere] = 3.0
    return u


def test_velocity_gradient_shape():
    """Check gradient output shape."""
    u = _synthetic_field(16, 1)
    grad = velocity_gradient(u[0])
    assert grad.shape == (16, 16, 16, 3, 3), f"Expected (16,16,16,3,3) got {grad.shape}"
    print("  ✓ velocity_gradient shape")


def test_scar_density():
    """Check scar density is non-negative."""
    u = _synthetic_field(16, 1)
    grad = velocity_gradient(u[0])
    mu = scar_density(grad)
    assert mu.shape == (16, 16, 16), f"Expected (16,16,16) got {mu.shape}"
    assert (mu >= 0).all(), "Scar density must be non-negative"
    print("  ✓ scar_density shape and non-negativity")


def test_threshold():
    """Check threshold produces binary mask."""
    mu = np.random.rand(16, 16, 16)
    mask = threshold_scar_support(mu, percentile=50.0)
    assert mask.dtype == np.uint8, f"Expected uint8 got {mask.dtype}"
    assert set(np.unique(mask)) <= {0, 1}, "Mask must be binary"
    frac = mask.mean()
    assert 0.45 < frac < 0.55, f"50th percentile mask fraction {frac} not ~0.5"
    print("  ✓ threshold_scar_support binary with correct fraction")


def test_betti_uniform():
    """Uniform field → β₂ = 0."""
    u = _synthetic_field(16, 1, void=False)
    grad = velocity_gradient(u[0])
    mu = scar_density(grad)
    mask = threshold_scar_support(mu, percentile=75)
    betti = compute_betti_numbers(mask, max_dim=2)
    assert betti.get(2, 0) == 0, f"β₂ should be 0, got {betti}"
    print("  ✓ β₂ = 0 for uniform field")


def test_betti_void():
    """Field with enclosed void → β₂ > 0."""
    u = _synthetic_field(16, 1, void=True)
    grad = velocity_gradient(u[0])
    mu = scar_density(grad)
    mask = threshold_scar_support(mu, percentile=75)
    betti = compute_betti_numbers(mask, max_dim=2)
    assert betti.get(2, 0) > 0, f"β₂ should be > 0 for void field, got {betti}"
    print(f"  ✓ β₂ = {betti.get(2)} > 0 for void field")


def test_pipeline_output():
    """Full pipeline produces correct JSON and PNG."""
    u = _synthetic_field(32, 5, void=True)
    with tempfile.TemporaryDirectory() as tmpdir:
        out = run_pipeline(u, percentile=75, outdir=tmpdir, prefix="test")
        assert len(out["timesteps"]) == 5
        assert out["summary"].startswith("β₂ spiked")
        # check files
        assert (Path(tmpdir) / "test.json").exists()
        assert (Path(tmpdir) / "test.png").exists()
        # verify JSON round-trip
        with open(Path(tmpdir) / "test.json") as f:
            loaded = json.load(f)
        assert loaded["shape"] == [5, 32, 32, 32, 3]
    print("  ✓ pipeline output correct")


def test_fallback_betti():
    """Ensure fallback works without gudhi."""
    try:
        import gudhi
    except ImportError:
        pass
    else:
        # Force fallback
        pass

    # Random mask — just check it returns something sane
    mask = np.zeros((16, 16, 16), dtype=np.uint8)
    mask[4:12, 4:12, 4:12] = 1  # solid cube → no voids
    betti = compute_betti_numbers(mask, max_dim=2)
    assert betti.get(2, 0) == 0, f"Solid cube should have β₂=0, got {betti}"
    print("  ✓ fallback: solid cube β₂ = 0")


if __name__ == "__main__":
    print("betti_tracker unit tests\n")
    test_velocity_gradient_shape()
    test_scar_density()
    test_threshold()
    test_betti_uniform()
    test_betti_void()
    test_fallback_betti()
    test_pipeline_output()
    print("\n✓ ALL TESTS PASSED")