#!/usr/bin/env python3
"""betti_tracker.py — NK-Hodge-FAMM Betti number tracker.

Computes Betti numbers (β₀, β₁, β₂) of the FAMM scar support
from a velocity field u(x,t).  β₂(scar support) > 0 is the
early-warning signal for Navier-Stokes blowup.

Usage:
  python3 betti_tracker.py field.npy [--percentile P] [--dx D] [--outdir O]
  python3 betti_tracker.py --test            # synthetic smoke test
"""

from __future__ import annotations
import argparse, json, os, sys, warnings
from pathlib import Path

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import ndimage as ndi

try:
    import gudhi as gd
    HAVE_GUDHI = True
except ImportError:
    HAVE_GUDHI = False

warnings.filterwarnings("ignore", category=UserWarning)


# ── helpers ──────────────────────────────────────────────────────────

def velocity_gradient(u: np.ndarray, dx: float = 1.0) -> np.ndarray:
    """∇u via central finite differences.

    Returns shape (..., 3, 3) where the last two axes are
    ∇u_ij = ∂u_i / ∂x_j.
    """
    ndim = u.ndim
    spatial_axes = tuple(range(ndim - 1))  # all axes except the last (velocity-component axis)
    grads = np.gradient(u, dx, axis=spatial_axes)
    return np.stack(grads, axis=-1)


def scar_density(grad_u: np.ndarray) -> np.ndarray:
    """μ = ‖∇u‖²_F — squared Frobenius norm of the (..., 3, 3) tensor."""
    return np.sum(grad_u ** 2, axis=(-1, -2))


def threshold_scar_support(mu: np.ndarray, percentile: float = 75.0) -> np.ndarray:
    """Binary mask where μ > δ (δ = percentile of μ at t=0 by default)."""
    delta = np.percentile(mu, percentile)
    return (mu > delta).astype(np.uint8)


# ── cubical persistence (gudhi) ─────────────────────────────────────

def _pad_cubical(value_array: np.ndarray, pad: int = 2):
    """Pad with background value (1.0) so gudhi sees no foreground at boundary."""
    return np.pad(value_array, pad, mode="constant", constant_values=1.0)


def compute_betti_numbers(binary_mask: np.ndarray, max_dim: int = 2):
    """β₀, β₁, β₂ of the foreground set via cubical persistence (gudhi).

    Convention: foreground value is LOW (appears at filtration 0),
    background is HIGH (added later).  Only essential (infinite-lifetime)
    features of the foreground are counted.

    Parameters
    ----------
    binary_mask : (N, N, N) uint8 — foreground = 1, background = 0.
    max_dim : int — maximum homology dimension to compute.

    Returns
    -------
    dict {dim: count}  e.g. {0: 1, 1: 0, 2: 1}
    """
    if not HAVE_GUDHI:
        return _betti_fallback(binary_mask, max_dim)

    # foreground = 0.0 (appears first), background = 1.0 (added later)
    cells = np.ones_like(binary_mask, dtype=np.float64)
    cells[binary_mask.astype(bool)] = 0.0
    padded = _pad_cubical(cells, pad=2)
    try:
        cc = gd.CubicalComplex(top_dimensional_cells=padded)
        cc.compute_persistence()
        dgms = cc.persistence()
    except Exception:
        return _betti_fallback(binary_mask, max_dim)

    betti: dict[int, int] = {d: 0 for d in range(max_dim + 1)}
    for dim, (birth, death) in dgms:
        if dim > max_dim or dim < 0:
            continue
        # Count features born at the start (foreground-only complex).
        # This includes both essential features (death = inf) and
        # features killed by background addition (death = 1.0).
        # Exclude zero-lifetime noise features.
        eps = 1e-9
        if birth < eps:
            lifetime = 1.0 if death == np.inf else (death - birth)
            if lifetime > eps:
                betti[dim] = betti.get(dim, 0) + 1
    return betti


# ── fallback: Euler characteristic + connected components ──────────

def _betti_fallback(binary_mask: np.ndarray, max_dim: int = 2):
    """Fallback Betti estimator w/ Euler characteristic + CC analysis.

    Uses:
        β₀ = foreground connected components
        β₂ = background CC − 1 (outer component)
        β₁ = β₀ + β₂ − χ  (Euler characteristic)
    """
    fg = binary_mask.astype(bool)
    if not fg.any():
        return {d: 0 for d in range(max_dim + 1)}

    # β₀
    labeled_fg, n_fg = ndi.label(fg)

    # Euler characteristic via voxel counts
    V = fg.sum()
    # edges (6-connectivity)
    edges = (
        (fg[:-1, :, :] & fg[1:, :, :]).sum() +
        (fg[:, :-1, :] & fg[:, 1:, :]).sum() +
        (fg[:, :, :-1] & fg[:, :, 1:]).sum()
    )
    # faces (4 corners forming a 2×2 square)
    faces = (
        (fg[:-1, :-1, :] & fg[1:, :-1, :] & fg[:-1, 1:, :] & fg[1:, 1:, :]).sum() +
        (fg[:-1, :, :-1] & fg[1:, :, :-1] & fg[:-1, :, 1:] & fg[1:, :, 1:]).sum() +
        (fg[:, :-1, :-1] & fg[:, 1:, :-1] & fg[:, :-1, 1:] & fg[:, 1:, 1:]).sum()
    )
    # cubes (8 vertices)
    cubes = (fg[:-1, :-1, :-1] & fg[1:, :-1, :-1] & fg[:-1, 1:, :-1] &
             fg[1:, 1:, :-1] & fg[:-1, :-1, 1:] & fg[1:, :-1, 1:] &
             fg[:-1, 1:, 1:] & fg[1:, 1:, 1:]).sum()

    chi = V - edges + faces - cubes

    # β₂ via complement CC
    bg = (~fg).astype(bool)
    labeled_bg, n_bg = ndi.label(bg)

    b0 = n_fg
    b2 = n_bg - 1 if n_bg > 0 else 0  # subtract outer component
    b1 = b0 + b2 - chi

    result = {0: b0}
    if max_dim >= 1:
        result[1] = max(b1, 0)
    if max_dim >= 2:
        result[2] = max(b2, 0)
    return result


# ── pipeline ─────────────────────────────────────────────────────────

def run_pipeline(
    u: np.ndarray,
    percentile: float = 75.0,
    dx: float = 1.0,
    outdir: str = ".",
    prefix: str = "betti",
) -> dict:
    """Full NK-Hodge-FAMM pipeline.

    Parameters
    ----------
    u : (T, N, N, N, 3)  velocity field.
    percentile : threshold percentile for scar support.
    dx : grid spacing.
    outdir : output directory.
    prefix : file prefix for JSON / PNG.

    Returns
    -------
    dict of results.
    """
    T = u.shape[0]
    results: list[dict] = []

    for t in range(T):
        grad = velocity_gradient(u[t], dx)
        mu = scar_density(grad)
        mask = threshold_scar_support(mu, percentile)
        betti = compute_betti_numbers(mask, max_dim=2)

        results.append({
            "t": int(t),
            "scar_density_mean": float(mu.mean()),
            "scar_density_max": float(mu.max()),
            "scar_support_fraction": float(mask.mean()),
            "betti_0": int(betti.get(0, 0)),
            "betti_1": int(betti.get(1, 0)),
            "betti_2": int(betti.get(2, 0)),
        })

    out_path = Path(outdir)
    out_path.mkdir(parents=True, exist_ok=True)

    # summary
    b2_series = [r["betti_2"] for r in results]
    spikes = [r["t"] for r in results if r["betti_2"] > 0]
    if spikes:
        summary = f"β₂ spiked at t={spikes}"
    else:
        summary = "β₂ stayed 0 throughout"

    output = {
        "shape": list(u.shape),
        "percentile": percentile,
        "dx": dx,
        "summary": summary,
        "timesteps": results,
    }

    # JSON
    json_path = out_path / f"{prefix}.json"
    with open(json_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"Wrote {json_path}")

    # plot
    fig, axes = plt.subplots(4, 1, figsize=(10, 10), sharex=True)

    ts = [r["t"] for r in results]
    b0 = [r["betti_0"] for r in results]
    b1 = [r["betti_1"] for r in results]
    b2 = [r["betti_2"] for r in results]

    axes[0].plot(ts, b0, "o-", color="C0")
    axes[0].set_ylabel("β₀")
    axes[0].set_title("Betti numbers vs time")

    axes[1].plot(ts, b1, "s-", color="C1")
    axes[1].set_ylabel("β₁")

    axes[2].plot(ts, b2, "D-", color="C2")
    axes[2].set_ylabel("β₂")
    axes[2].axhline(y=0, color="gray", linestyle="--", alpha=0.5)

    scar_frac = [r["scar_support_fraction"] for r in results]
    axes[3].plot(ts, scar_frac, ".-", color="C3")
    axes[3].set_ylabel("scar fraction")
    axes[3].set_xlabel("time step")

    fig.suptitle(summary, fontsize=11)
    plt.tight_layout()
    png_path = out_path / f"{prefix}.png"
    fig.savefig(png_path, dpi=150)
    plt.close(fig)
    print(f"Wrote {png_path}")

    print(f"\nSummary: {summary}")
    return output


# ── smoke test ───────────────────────────────────────────────────────

def _synthetic_field(shape: tuple[int, ...], void: bool = False) -> np.ndarray:
    """Create a synthetic (T, N, N, N, 3) velocity field.

    Uniform case: constant velocity everywhere → zero gradients →
    empty scar support → β₂ = 0.

    Void case: uniform background with a spherical region of different
    velocity → sharp gradient shell at boundary → scar support forms
    a spherical shell enclosing a void → β₂ > 0.
    """
    T, N = shape[0], shape[1]
    u = np.ones(shape, dtype=np.float64) * 0.5  # constant → zero gradient

    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  # sharp contrast → gradient boundary → void
    return u


def _test():
    N = 32
    T = 10
    print("=" * 56)
    print("betti_tracker — smoke test")
    print("=" * 56)

    # Test 1: uniform field → β₂ = 0
    print("\n[Test 1] Uniform field (expect β₂ = 0)")
    u_flat = _synthetic_field((T, N, N, N, 3), void=False)
    out1 = run_pipeline(u_flat, percentile=75, outdir="/tmp/betti_test_1", prefix="flat")
    assert all(r["betti_2"] == 0 for r in out1["timesteps"]), "FAIL: β₂ should be 0"
    print("  ✓ β₂ = 0 at all timesteps")

    # Test 2: field with enclosed void → β₂ > 0
    print("\n[Test 2] Field with spherical void (expect β₂ > 0)")
    u_void = _synthetic_field((T, N, N, N, 3), void=True)
    out2 = run_pipeline(u_void, percentile=75, outdir="/tmp/betti_test_2", prefix="void")
    b2_series = [r["betti_2"] for r in out2["timesteps"]]
    assert any(v > 0 for v in b2_series), f"FAIL: β₂ should be > 0 somewhere, got {b2_series}"
    print(f"  ✓ β₂ > 0 at t={[r['t'] for r in out2['timesteps'] if r['betti_2'] > 0]}")

    print("\n" + "=" * 56)
    print("ALL TESTS PASSED")
    print("=" * 56)


# ── CLI ──────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(
        description="NK-Hodge-FAMM Betti number tracker"
    )
    ap.add_argument("field", nargs="?", help="Path to .npy velocity field (T,N,N,N,3)")
    ap.add_argument("--percentile", type=float, default=75.0, help="Scar threshold percentile (default 75)")
    ap.add_argument("--dx", type=float, default=1.0, help="Grid spacing (default 1.0)")
    ap.add_argument("--outdir", default=".", help="Output directory")
    ap.add_argument("--prefix", default="betti", help="Output file prefix")
    ap.add_argument("--test", action="store_true", help="Run smoke test")
    args = ap.parse_args()

    if args.test:
        _test()
        return

    if args.field is None:
        ap.print_help()
        sys.exit(1)

    u = np.load(args.field)
    run_pipeline(u, percentile=args.percentile, dx=args.dx,
                 outdir=args.outdir, prefix=args.prefix)


if __name__ == "__main__":
    main()