#!/usr/bin/env python3
"""taylor_green_betti.py — Betti tracking on Taylor-Green vortex.

Tests whether β₂ of FAMM scar support correctly identifies enclosed
voids in a Navier-Stokes-like velocity field.
"""

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

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

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

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


# ── Field generators ─────────────────────────────────────────────────

def taylor_green_field(
    N: int = 32,
    nu: float = 0.1,
    t: float = 0.0,
    void_center: tuple[float, float, float] | None = None,
    void_radius: float = 0.0,
    void_velocity: tuple[float, float, float] | None = None,
) -> np.ndarray:
    """Taylor-Green vortex: (N,N,N,3) velocity field.

    u = sin(x) cos(y) cos(z) e^{-2νt}
    v = -cos(x) sin(y) cos(z) e^{-2νt}
    w = 0

    Void parameters use grid-index coordinates.
    """
    x = np.linspace(0, 2 * np.pi, N, endpoint=False)
    y = np.linspace(0, 2 * np.pi, N, endpoint=False)
    z = np.linspace(0, 2 * np.pi, N, endpoint=False)
    X, Y, Z = np.meshgrid(x, y, z, indexing="ij")

    decay = np.exp(-2 * nu * t)
    u = np.zeros((N, N, N, 3), dtype=np.float64)
    u[..., 0] = np.sin(X) * np.cos(Y) * np.cos(Z) * decay
    u[..., 1] = -np.cos(X) * np.sin(Y) * np.cos(Z) * decay

    if void_center is not None and void_radius > 0:
        cx, cy, cz = void_center
        ix, iy, iz = np.meshgrid(np.arange(N), np.arange(N), np.arange(N),
                                 indexing="ij")
        dist2 = (ix - cx) ** 2 + (iy - cy) ** 2 + (iz - cz) ** 2
        sphere = dist2 <= void_radius ** 2
        if void_velocity is not None:
            u[sphere] = np.array(void_velocity, dtype=np.float64)
        else:
            u[sphere] = 0.0
    return u


def uniform_field(
    N: int = 32,
    void_center: tuple[float, float, float] | None = None,
    void_radius: float = 0.0,
) -> np.ndarray:
    """Constant velocity field with optional spherical void."""
    u = np.ones((N, N, N, 3), dtype=np.float64) * 0.5
    if void_center is not None and void_radius > 0:
        cx, cy, cz = void_center
        ix, iy, iz = np.meshgrid(np.arange(N), np.arange(N), np.arange(N),
                                 indexing="ij")
        dist2 = (ix - cx) ** 2 + (iy - cy) ** 2 + (iz - cz) ** 2
        sphere = dist2 <= void_radius ** 2
        u[sphere] = 3.0
    return u


def noise_field(N: int = 32, seed: int = 42) -> np.ndarray:
    """Random velocity field (uniform in [-1, 1])."""
    return np.random.default_rng(seed).uniform(-1.0, 1.0, (N, N, N, 3))


# ── Analysis ─────────────────────────────────────────────────────────

def analyze_field(u: np.ndarray, percentile: float = 75.0,
                  dx: float = 1.0) -> dict:
    """Run NK-Hodge-FAMM pipeline on a single velocity field."""
    grad = velocity_gradient(u, dx)
    mu = scar_density(grad)
    mask = threshold_scar_support(mu, percentile)
    betti_raw = compute_betti_numbers(mask, max_dim=2)
    betti = {int(k): int(v) for k, v in betti_raw.items()}
    return {
        "mask": mask,
        "mu": mu,
        "betti": betti,
        "scar_density_mean": float(mu.mean()),
        "scar_density_max": float(mu.max()),
        "scar_support_fraction": float(mask.mean()),
    }


# ── Plotting ─────────────────────────────────────────────────────────

def make_plot(results: list[dict], names: list[str], save_path: str):
    """Layout: 2×3 scar mid-slices, then bar+table+interpretation, then histogram."""
    n = len(results)
    fig = plt.figure(figsize=(14, 14))
    gs = fig.add_gridspec(4, 3, hspace=0.35, wspace=0.25)

    # Rows 0-1: scar support mid-slices (2 rows × 3 cols)
    img_axes = []
    for r in range(2):
        row = []
        for c in range(3):
            ax = fig.add_subplot(gs[r, c])
            row.append(ax)
        img_axes.append(row)

    for i, (res, name) in enumerate(zip(results, names)):
        r, c = divmod(i, 3)
        ax = img_axes[r][c]
        mid = res["mask"].shape[2] // 2
        ax.imshow(res["mask"][:, :, mid], cmap="Reds", interpolation="nearest")
        b = res["betti"]
        ax.set_title(f"{name}\nβ₀={b[0]}  β₁={b[1]}  β₂={b[2]}", fontsize=9)
        ax.axis("off")

    # Row 2: bar chart (col 0), table (col 1), interpretation (col 2)
    ax_bar = fig.add_subplot(gs[2, 0])
    ax_tab = fig.add_subplot(gs[2, 1])
    ax_int = fig.add_subplot(gs[2, 2])

    labels = names
    x = np.arange(n)
    w = 0.25
    ax_bar.bar(x - w, [r["betti"][0] for r in results], w, label="β₀", color="C0")
    ax_bar.bar(x, [r["betti"][1] for r in results], w, label="β₁", color="C1")
    ax_bar.bar(x + w, [r["betti"][2] for r in results], w, label="β₂", color="C2")
    ax_bar.set_xticks(x)
    ax_bar.set_xticklabels(labels, fontsize=8)
    ax_bar.set_ylabel("Betti number")
    ax_bar.legend(fontsize=8)
    ax_bar.set_title("Betti numbers")

    cell_text = [[str(r["betti"][d]) for d in (0, 1, 2)] for r in results]
    ax_tab.axis("off")
    tbl = ax_tab.table(cellText=cell_text, rowLabels=labels,
                        colLabels=["β₀", "β₁", "β₂"],
                        loc="center", cellLoc="center")
    tbl.auto_set_font_size(False)
    tbl.set_fontsize(8)
    tbl.scale(1, 1.4)
    ax_tab.set_title("Betti table")

    lines = ["β₂ Interpretation:", "─" * 30]
    for name, r in zip(names, results):
        b2 = r["betti"][2]
        lines.append(f"  {name:>15s}: β₂={b2}  → {'voids' if b2 > 0 else 'no voids'}")
    lines.append("")
    lines.append("Void detection:")
    lines.append("  β₂>0 ⇔ scar support encloses cavity")
    lines.append("  Uniform+void: clean shell → β₂>0")
    lines.append("  TG+void: depends on void vs percolation")
    ax_int.axis("off")
    ax_int.text(0, 1.0, "\n".join(lines), fontsize=9,
                verticalalignment="top", family="monospace")
    ax_int.set_title("Interpretation")

    # Row 3: histogram (full width)
    ax_det = fig.add_subplot(gs[3, :])
    mu_flat = np.concatenate([r["mu"].ravel() for r in results])
    ax_det.hist(mu_flat, bins=80, density=True, alpha=0.4, color="gray",
                label="all scenarios")
    for i, (r, name) in enumerate(zip(results, names)):
        ax_det.axvline(x=r["scar_density_max"], color=f"C{i}", ls="--",
                       lw=1, label=f"{name} σ_max={r['scar_density_max']:.4f}")
    ax_det.set_xlabel("scar density μ(x)")
    ax_det.set_ylabel("density")
    ax_det.set_title("Scar density distributions")
    ax_det.legend(fontsize=6)

    fig.suptitle("Betti Tracker — Taylor-Green Vortex Analysis", fontsize=13)
    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"Wrote {save_path}")


# ── Scenarios ─────────────────────────────────────────────────────────

N = 32
VOID_CENTER = (24, 16, 16)  # off-center where TG velocity ≠ 0

SCENARIOS = [
    ("Uniform flow",
     lambda: uniform_field(N=N)),
    ("Uniform+void",
     lambda: uniform_field(N=N, void_center=VOID_CENTER, void_radius=5)),
    ("TG smooth",
     lambda: taylor_green_field(N=N, t=0.0)),
    ("TG+void r=5",
     lambda: taylor_green_field(N=N, t=0.0, void_center=VOID_CENTER,
                                void_radius=5)),
    ("TG+void r=6",
     lambda: taylor_green_field(N=N, t=0.0, void_center=VOID_CENTER,
                                void_radius=6)),
    ("Noise",
     lambda: noise_field(N=N)),
]


def main():
    outdir = Path(__file__).resolve().parent
    results, entries = [], []

    print("=" * 60)
    print("Betti Tracker — Taylor-Green Vortex Analysis")
    print("=" * 60)

    for name, gen_fn in SCENARIOS:
        print(f"\n  [{name}]")
        u = gen_fn()
        res = analyze_field(u, percentile=75.0)
        b = res["betti"]
        b2 = b[2]

        interp = f"β₂={b2} — indicating voids in scar support" if b2 > 0 \
            else f"β₂={b2} — no voids in scar support"
        print(f"    β₀={b[0]}  β₁={b[1]}  β₂={b2}")
        print(f"    scar fraction={res['scar_support_fraction']:.3f}")
        print(f"    {interp}")

        entries.append({
            "scenario": name,
            "betti_0": b[0], "betti_1": b[1], "betti_2": b2,
            "interpretation": interp,
            "scar_density_mean": res["scar_density_mean"],
            "scar_density_max": res["scar_density_max"],
            "scar_support_fraction": res["scar_support_fraction"],
        })
        results.append(res)

    report = dict(
        schema="taylor_green_betti_v1",
        description="Betti number analysis of Taylor-Green vortex with voids",
        grid=f"{N}x{N}x{N}",
        percentile=75.0,
        scenarios=entries,
        summary={entries[i]["scenario"]: entries[i]["betti_2"]
                 for i in range(len(entries))},
    )
    json_path = outdir / "taylor_green_betti.json"
    with open(json_path, "w") as f:
        json.dump(report, f, indent=2)
    print(f"\nWrote {json_path}")

    make_plot(results, [s[0] for s in SCENARIOS],
              str(outdir / "taylor_green_betti.png"))

    print("\n" + "=" * 60)
    print("DONE")
    print("=" * 60)


# ── Entry point ───────────────────────────────────────────────────────

if __name__ == "__main__":
    test = "--test" in sys.argv[1:]
    if test:
        sys.argv.remove("--test")
    main()
    if test:
        with open(Path(__file__).resolve().parent / "taylor_green_betti.json") as f:
            rpt = json.load(f)
        s = rpt["summary"]
        ok = True
        if s["Uniform flow"] != 0:
            print(f"\nFAIL: uniform flow expected β₂=0, got β₂={s['Uniform flow']}")
            ok = False
        if s["Uniform+void"] <= 0:
            print(f"\nFAIL: uniform+void expected β₂>0, got β₂={s['Uniform+void']}")
            ok = False
        if s["TG smooth"] != 0:
            print(f"\nWARN: smooth TG has β₂={s['TG smooth']} (expected 0)")
        if s["TG+void r=5"] <= 0:
            print(f"\nWARN: TG+void r=5 has β₂={s['TG+void r=5']} (expected >0)")
        else:
            print(f"\n✓ TG+void r=5: β₂={s['TG+void r=5']} > 0 — void detected")
        if ok:
            print("✓ ALL HARD EXPECTATIONS MET")
        else:
            print("✗ SOME HARD EXPECTATIONS FAILED")
            sys.exit(1)