"""
Next-Generation Curvature Parameterization

Implements four major upgrades to probe mass-curvature relationships:
- Upgrade A: Exponential generators (non-linear deformations)
- Upgrade B: Holonomy-space parameterization (gauge field geometry)
- Upgrade C: Spinor states (full Clifford DOF)
- Upgrade D: Time-dependent QGT (transient curvature)

This module tests the hypothesis that mass emerges from curvature in
gauge field coupling-space or from transient curvature events during
geometric relaxation.
"""

import jax
import jax.numpy as jnp
from jax import jit
import sys
import os
from typing import Dict, Any, Optional, List, Tuple
from datetime import datetime
from pathlib import Path
import pandas as pd
import numpy as np

# Add project root
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))

from light_theory_realm.engine import CliffordEngine
from light_theory_realm.qgt import CliffordQGT
from light_theory_realm.experiments.prime_gauge.prime_plaquette import PrimeWilsonLoop
from light_theory_realm.experiments.prime_gauge.prime_sixth_order_flow import (
    build_prime_ring,
    build_sixth_order_operator,
    gradient_flow_step,
    compute_S_rep,
    compute_local_S_rep,
    compute_gradient_norm,
)
from light_theory_realm.pocket_u_lite.plaquettes import (
    PLAQUETTES,
    PDG_MASSES,
    ALL_PARTICLES,
)


# =============================================================================
# Upgrade A: Exponential Generator Parameterization
# =============================================================================

def build_edge_generators(
    primes: List[int],
    engine: CliffordEngine,
) -> List[jnp.ndarray]:
    """
    Build edge generators from Wilson line operators.
    
    Args:
        primes: Prime sequence [p₁, p₂, p₃, p₄]
        engine: CliffordEngine instance
    
    Returns:
        List of 4 edge generators [X₁, X₂, X₃, X₄]
        Each Xᵢ is (4, 4) matrix acting on node space
    
    Construction:
        For each edge i→j:
        - Compute phase: φ_ij = ln(p_j/p_i)
        - Build generator: X_ij = φ_ij * (node_i - node_j) operator
    """
    n_nodes = len(primes)
    generators = []
    
    for i in range(n_nodes):
        j = (i + 1) % n_nodes
        
        # Phase from prime ratio
        phi_ij = np.log(primes[j] / primes[i])
        
        # Build edge operator: shifts amplitude from node i to node j
        X_ij = np.zeros((n_nodes, n_nodes))
        X_ij[j, i] = phi_ij  # Forward edge
        X_ij[i, j] = -phi_ij  # Backward edge (anti-symmetric)
        
        generators.append(jnp.array(X_ij))
    
    return generators


def project_generators_to_dim(generators: List[jnp.ndarray], n: int) -> List[jnp.ndarray]:
    """
    Reduce a list of generator matrices to act on an n×n node subspace.
    
    If the original generators are (N,N) with N>n, take the top-left n×n block.
    This avoids shape mismatches when analyzing 2D curvature slices.
    
    Args:
        generators: List of (N,N) generator matrices
        n: Target dimension for projection
    
    Returns:
        List of n generators, each (n,n)
    """
    projected = []
    for G in generators[:n]:
        projected.append(G[:n, :n])
    return projected


def build_exponential_state(
    theta: jnp.ndarray,
    generators: List[jnp.ndarray],
    base_psi: jnp.ndarray,
) -> jnp.ndarray:
    """
    Build state via exponential map: ψ(θ) = exp(Σ θᵃ Xₐ) ψ₀
    
    Implemented with a **6th-order Taylor expansion** of the matrix exponential
    to stay faithful to the sixth-order structure of the theory, while remaining
    JAX-differentiable.
    
    The sixth-order truncation matches the order of the Master Information 
    Equation (Δ³), preserving theoretical coherence.
    
    Args:
        theta: Parameters (n_params,)
        generators: List of generator matrices
        base_psi: Base state
    
    Returns:
        Parameterized state normalized on projective Hilbert space
    """
    # Build Lie algebra element: X_total = Σ θᵃ Xₐ
    X_total = sum(theta[a] * generators[a] for a in range(len(theta)))

    # Sixth-order Taylor expansion of exp(X_total)
    # exp(X) = I + X + X²/2! + X³/3! + X⁴/4! + X⁵/5! + X⁶/6! + O(X⁷)
    n = X_total.shape[0]
    I = jnp.eye(n, dtype=X_total.dtype)

    X2 = X_total @ X_total
    X3 = X2 @ X_total
    X4 = X3 @ X_total
    X5 = X4 @ X_total
    X6 = X5 @ X_total

    U = (
        I
        + X_total
        + X2 / 2.0
        + X3 / 6.0           # 3! = 6
        + X4 / 24.0          # 4! = 24
        + X5 / 120.0         # 5! = 120
        + X6 / 720.0         # 6! = 720
    )

    # Apply to base state
    psi = U @ base_psi

    # Normalize to stay on projective Hilbert space
    norm = jnp.linalg.norm(psi) + 1e-12
    return psi / norm


# =============================================================================
# Option B: Full 2D Ricci Scalar Curvature from QGT Metric
# =============================================================================

def metric_from_theta(
    theta: jnp.ndarray,
    primes: List[int],
    engine: CliffordEngine,
    base_psi: jnp.ndarray,
    generators: List[jnp.ndarray],
    qgt: CliffordQGT,
) -> jnp.ndarray:
    """
    Given θ, build ψ(θ), its Jacobian, and return the Fisher metric g_ab(θ).
    
    This is the 2×2 information metric on parameter space, computed from
    the QGT real part (Fisher information).
    
    Args:
        theta: Parameter point (2,)
        primes: Prime sequence
        engine: CliffordEngine instance
        base_psi: Base state for exponential map
        generators: List of generator matrices (2,)
        qgt: CliffordQGT instance
    
    Returns:
        Fisher metric g_ab(θ) as (2,2) symmetric matrix
    """
    def psi_fn(th):
        return build_exponential_state(th, generators, base_psi)
    
    psi = psi_fn(theta)
    # dψ/dθᵃ: shape (n_spinor, n_params)
    jacobian = jax.jacfwd(psi_fn)(theta)
    
    # Reshape for QGT: (n_params, n_spinor)
    jac_reshaped = jacobian.T
    
    # Add complex phase for non-trivial geometry
    psi_complex = psi + 1j * 0.1 * jnp.arange(len(psi)) / len(psi)
    psi_complex = psi_complex / jnp.linalg.norm(psi_complex)
    psi_spinor = psi_complex[:, None]
    
    fisher, _ = qgt.compute_full_qgt(psi_spinor, jac_reshaped)
    
    # Ensure Hermitian/symmetric numerical stability
    fisher_sym = 0.5 * (fisher + fisher.T)
    return jnp.real(fisher_sym)


def christoffel_symbols_2d(
    theta: jnp.ndarray,
    primes: List[int],
    engine: CliffordEngine,
    base_psi: jnp.ndarray,
    generators: List[jnp.ndarray],
    qgt: CliffordQGT,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """
    Compute metric g_ab(θ), its inverse, and Γ^a_bc(θ) for a,b,c in {0,1}.
    
    Uses JAX autodiff to compute metric derivatives ∂_c g_ab, then constructs
    Christoffel symbols via the standard formula:
    
    Γ^a_bc = ½ g^ad (∂_b g_dc + ∂_c g_db - ∂_d g_bc)
    
    Args:
        theta: Parameter point (2,)
        primes: Prime sequence
        engine: CliffordEngine instance
        base_psi: Base state for exponential map
        generators: List of generator matrices (2,)
        qgt: CliffordQGT instance
    
    Returns:
        g      : (2,2) metric
        g_inv  : (2,2) inverse metric
        Gamma  : (2,2,2) Christoffel symbols with indices [a,b,c] = Γ^a_bc
    """
    # Metric
    g = metric_from_theta(theta, primes, engine, base_psi, generators, qgt)
    g_reg = g + 1e-12 * jnp.eye(2)
    g_inv = jnp.linalg.inv(g_reg)

    # ∂_c g_ab via jacobian over theta
    # dg[c, a, b] = ∂_c g_ab
    def g_fn(th):
        return metric_from_theta(th, primes, engine, base_psi, generators, qgt)

    dg = jax.jacfwd(g_fn)(theta)  # shape (2, 2, 2)

    # Gamma[a,b,c] = Γ^a_bc
    # Using explicit computation for clarity (2×2×2 is tiny)
    Gamma = jnp.zeros((2, 2, 2), dtype=g.dtype)
    
    for a in range(2):
        for b in range(2):
            for c in range(2):
                # Γ^a_bc = ½ Σ_d g^ad (∂_b g_dc + ∂_c g_db - ∂_d g_bc)
                term = 0.0
                for d in range(2):
                    term += g_inv[a, d] * (
                        dg[b, d, c]      # ∂_b g_dc
                        + dg[c, d, b]    # ∂_c g_db
                        - dg[d, b, c]    # ∂_d g_bc
                    )
                Gamma = Gamma.at[a, b, c].set(0.5 * term)

    return g, g_inv, Gamma


def scalar_curvature_full_2d(
    theta: jnp.ndarray,
    primes: List[int],
    engine: CliffordEngine,
    base_psi: jnp.ndarray,
    generators: List[jnp.ndarray],
    qgt: CliffordQGT,
) -> Tuple[float, jnp.ndarray, jnp.ndarray]:
    """
    Compute the full 2D scalar curvature R(θ) from the QGT metric.
    
    This is Option B: no shortcuts, full differential geometry in 2D.
    
    Implements the complete Riemann curvature computation:
    1. Compute Christoffel symbols Γ^a_bc
    2. Compute derivatives ∂_e Γ^a_bc
    3. Compute Ricci tensor: R_bd = ∂_a Γ^a_bd - ∂_d Γ^a_ba + Γ^a_bd Γ^c_ac - Γ^a_bc Γ^c_ad
    4. Contract to scalar: R = g^bd R_bd
    
    Args:
        theta: Parameter point (2,)
        primes: Prime sequence
        engine: CliffordEngine instance
        base_psi: Base state for exponential map
        generators: List of generator matrices (2,)
        qgt: CliffordQGT instance
    
    Returns:
        R      : scalar curvature
        g      : metric (2,2)
        Gamma  : Christoffel symbols (2,2,2)
    """
    g, g_inv, Gamma = christoffel_symbols_2d(theta, primes, engine, base_psi, generators, qgt)

    # We need ∂_e Γ^a_bc for e in {0,1}
    def Gamma_fn(th):
        _, _, G = christoffel_symbols_2d(th, primes, engine, base_psi, generators, qgt)
        return G

    # dGamma[e, a, b, c] = ∂_e Γ^a_bc
    dGamma = jax.jacfwd(Gamma_fn)(theta)

    # Ricci tensor R_bd
    Ric = jnp.zeros((2, 2), dtype=g.dtype)

    for b in range(2):
        for d in range(2):
            # ∂_a Γ^a_bd (sum over a)
            term1 = 0.0
            for a in range(2):
                term1 += dGamma[a, a, b, d]
            
            # ∂_d Γ^a_ba (sum over a)
            term2 = 0.0
            for a in range(2):
                term2 += dGamma[d, a, b, a]

            # Γ^a_bd Γ^c_ac (sum over a, c)
            term3 = 0.0
            for a in range(2):
                for c in range(2):
                    term3 += Gamma[a, b, d] * Gamma[c, a, c]
            
            # Γ^a_bc Γ^c_ad (sum over a, c)
            term4 = 0.0
            for a in range(2):
                for c in range(2):
                    term4 += Gamma[a, b, c] * Gamma[c, a, d]

            Ric = Ric.at[b, d].set(term1 - term2 + term3 - term4)

    # Scalar curvature R = g^bd R_bd
    R = jnp.sum(g_inv * Ric)
    
    return float(R), g, Gamma


def compute_qgt_exponential(
    theta: jnp.ndarray,
    generators: List[jnp.ndarray],
    base_psi: jnp.ndarray,
    engine: CliffordEngine,
    primes: List[int],
) -> Dict[str, Any]:
    """
    Compute QGT for exponential parameterization.
    
    Uses automatic differentiation to compute Jacobian via the 6th-order
    exponential map. Returns Fisher metric and Berry curvature with a
    principled Ricci scalar that captures manifold curvature.
    
    Returns:
        Dictionary with fisher, berry, curvature_scalar
    """
    # Build state via 6th-order exponential
    psi = build_exponential_state(theta, generators, base_psi)
    
    # Compute Jacobian via autodiff (respects 6th-order structure)
    def state_fn(th):
        return build_exponential_state(th, generators, base_psi)
    
    jacobian = jax.jacfwd(state_fn)(theta)  # shape (n_psi, n_params)
    
    # Transpose jacobian to match QGT expectation: (n_params, n_psi)
    jacobian = jacobian.T
    
    # Compute QGT: Fisher metric + Berry curvature
    qgt = CliffordQGT(engine)
    
    # Break symmetry with complex phase to reveal geometric structure
    psi_with_phase = psi + 1j * 0.1 * jnp.arange(len(psi)) / len(psi)
    psi_with_phase = psi_with_phase / jnp.linalg.norm(psi_with_phase)
    psi_spinor = psi_with_phase[:, None]
    
    # Compute Fisher (information metric) and Berry (topological curvature)
    fisher, berry = qgt.compute_full_qgt(psi_spinor, jacobian)
    
    # Option B: Full 2D Ricci Scalar from QGT Metric
    # Use only first 2 parameters for well-defined 2D curvature computation
    theta_2d = theta[:2]
    
    # Project generators to 2×2 subspace (not just slice - need 2×2 matrices!)
    generators_2d = project_generators_to_dim(generators, 2)
    
    # Rebuild base_psi for 2D (2 nodes, not slicing 4-element array)
    primes_2d = primes[:2]
    base_psi_2d = jnp.log(jnp.array(primes_2d, dtype=float))
    base_psi_2d = base_psi_2d / jnp.linalg.norm(base_psi_2d)
    
    # Compute full 2D scalar curvature via differential geometry
    R, fisher_2d, Gamma = scalar_curvature_full_2d(
        theta=theta_2d,
        primes=primes,
        engine=engine,
        base_psi=base_psi_2d,
        generators=generators_2d,
        qgt=qgt,
    )
    
    return {
        "psi": psi,
        "jacobian": jacobian,
        "fisher": fisher,
        "berry": berry,
        "curvature_scalar": R,
        "Gamma": Gamma,  # Christoffel symbols for diagnostics
    }


# =============================================================================
# Upgrade B: Holonomy-Space Parameterization
# =============================================================================

def build_parameterized_wilson_loop(
    primes: List[int],
    alpha: jnp.ndarray,
    engine: CliffordEngine,
) -> jnp.ndarray:
    """
    Build Wilson loop with parameterized couplings:
    W(α) = ∏ exp(i·α_ij·ln(p_j/p_i)·G_ij)
    
    Args:
        primes: Prime sequence [p₁, p₂, p₃, p₄]
        alpha: Coupling parameters (4,) - one per edge
        engine: CliffordEngine
    
    Returns:
        Wilson loop operator W(α)
    """
    qc_loop = PrimeWilsonLoop(engine)
    
    # Build loop with scaled phases
    n_edges = len(primes) - 1  # Assuming closed loop
    W = jnp.eye(4, dtype=jnp.complex64)  # Start with identity
    
    for i in range(n_edges):
        j = i + 1
        
        # Compute scaled phase
        phi_base = np.log(primes[j] / primes[i])
        phi_scaled = alpha[i] * phi_base
        
        # Build edge operator with scaled phase
        # Simplified: use phase rotation matrix
        edge_op = jnp.array([
            [jnp.cos(phi_scaled), -jnp.sin(phi_scaled), 0, 0],
            [jnp.sin(phi_scaled), jnp.cos(phi_scaled), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ], dtype=jnp.complex64)
        
        W = W @ edge_op
    
    return W


def compute_ground_state_of_wilson_hamiltonian(
    W_alpha: jnp.ndarray,
) -> jnp.ndarray:
    """
    Compute ground state of Hamiltonian built from Wilson loop.
    
    H(α) = -Re(W(α)) (simplified)
    ψ(α) = lowest eigenstate of H(α)
    
    Args:
        W_alpha: Wilson loop operator
    
    Returns:
        Ground state ψ(α)
    """
    # Build Hamiltonian from Wilson loop
    H_alpha = -jnp.real(W_alpha)
    
    # Find lowest eigenstate
    eigenvalues, eigenvectors = jnp.linalg.eigh(H_alpha)
    
    # Ground state = eigenvector with smallest eigenvalue
    psi_alpha = eigenvectors[:, 0]
    
    # Normalize
    return psi_alpha / jnp.linalg.norm(psi_alpha)


def compute_qgt_holonomy_space(
    alpha: jnp.ndarray,
    primes: List[int],
    engine: CliffordEngine,
) -> Dict[str, Any]:
    """
    Compute QGT in holonomy-space (α-space).
    
    Args:
        alpha: Coupling parameters (4,)
        primes: Prime sequence
        engine: CliffordEngine
    
    Returns:
        Dictionary with fisher, berry, curvature_scalar
    """
    # Build Wilson loop
    W_alpha = build_parameterized_wilson_loop(primes, alpha, engine)
    
    # Get ground state
    psi_alpha = compute_ground_state_of_wilson_hamiltonian(W_alpha)
    
    # Compute Jacobian: ∂ψ/∂α
    def state_fn(a):
        W = build_parameterized_wilson_loop(primes, a, engine)
        return compute_ground_state_of_wilson_hamiltonian(W)
    
    jacobian = jax.jacfwd(state_fn)(alpha)  # (n_psi, n_alpha)
    
    # Compute QGT
    qgt = CliffordQGT(engine)
    psi_spinor = psi_alpha[:, None]
    fisher, berry = qgt.compute_full_qgt(psi_spinor, jacobian)
    
    # Ricci scalar
    g_inv = jnp.linalg.inv(fisher + 1e-12 * jnp.eye(len(alpha)))
    R = -2.0 * jnp.trace(berry @ g_inv)
    
    return {
        "psi": psi_alpha,
        "W": W_alpha,
        "jacobian": jacobian,
        "fisher": fisher,
        "berry": berry,
        "curvature_scalar": float(R),
    }


# =============================================================================
# Upgrade C: Spinor State Extension
# =============================================================================

def build_spinor_state(
    engine: CliffordEngine,
    n_nodes: int = 4,
    seed: int = 42,
) -> jnp.ndarray:
    """
    Build spinor state: ψ ∈ Cl(1,3)^n_nodes
    
    Returns:
        Spinor state (n_nodes * 4,) - flattened spinor per node
    """
    key = jax.random.PRNGKey(seed)
    
    # Each node gets a 4-component Clifford spinor
    psi_spinor = jax.random.normal(key, (n_nodes, 4))
    
    # Flatten and normalize
    psi_flat = psi_spinor.flatten()
    return psi_flat / jnp.linalg.norm(psi_flat)


def extend_operator_to_spinor_space(
    H_scalar: jnp.ndarray,
    n_spinor_components: int = 4,
) -> jnp.ndarray:
    """
    Extend scalar operator H to act on spinor space.
    
    H_spinor = H_scalar ⊗ I₄
    """
    return jnp.kron(H_scalar, jnp.eye(n_spinor_components))


# =============================================================================
# Upgrade D: Time-Dependent QGT Tracking
# =============================================================================

def track_curvature_during_flow(
    primes: List[int],
    engine: CliffordEngine,
    mode: str = "exponential",  # "exponential" or "holonomy"
    n_steps: int = 100,
    dt: float = 0.01,
    sample_every: int = 10,
    verbose: bool = True,
) -> pd.DataFrame:
    """
    Track R(τ) and ρ_IGBP(τ) during gradient flow evolution.
    
    Args:
        primes: Prime sequence
        engine: CliffordEngine
        mode: "exponential" or "holonomy" parameterization
        n_steps: Evolution steps
        dt: Time step
        sample_every: Compute QGT every N steps
        verbose: Print progress
    
    Returns:
        DataFrame with [step, tau, R, rho_igbp, S_rep, ...]
    """
    # Build infrastructure
    _, laplacian, _ = build_prime_ring()
    L6, H = build_sixth_order_operator(laplacian)
    
    # Initialize state
    key = jax.random.PRNGKey(hash(str(primes)) % 2**32)
    psi_init = jax.random.normal(key, (4,))
    psi_init = psi_init / jnp.linalg.norm(psi_init)
    
    psi = psi_init
    history = []
    
    # Build generators/parameters based on mode
    if mode == "exponential":
        generators = build_edge_generators(primes, engine)
        # FIX #7: Use asymmetric base state from prime logs, not uniform
        base_psi = jnp.log(jnp.array(primes, dtype=float))
        base_psi = base_psi / jnp.linalg.norm(base_psi)
        
        # FIX #4: Initialize alpha for dynamic evolution
        alpha = jnp.ones(4)
    elif mode == "holonomy":
        # Will use alpha parameters
        alpha = jnp.ones(4)
    
    if verbose:
        print(f"Tracking curvature for {primes} (mode={mode})")
    
    for step in range(n_steps + 1):
        tau = step * dt
        
        # FIX #4: Evolve alpha dynamically to create transient curvature events
        if mode == "exponential" and step > 0:
            # Simple gradient step on alpha to maximize curvature
            # (This is a proxy for entropy-driven evolution)
            def alpha_loss(alph):
                theta_upd = fit_parameters_to_state(psi, generators, base_psi)
                qgt_res = compute_qgt_exponential(theta_upd, generators, base_psi, engine, primes)
                return -jnp.abs(qgt_res["curvature_scalar"])  # Maximize |R|
            
            alpha_grad = jax.grad(alpha_loss)(alpha)
            alpha = alpha - 0.001 * alpha_grad
        
        # Sample QGT if at sampling point
        if step % sample_every == 0:
            try:
                if mode == "exponential":
                    # Fit parameters to current state
                    theta_fit = fit_parameters_to_state(psi, generators, base_psi)
                    qgt_result = compute_qgt_exponential(theta_fit, generators, base_psi, engine, primes)
                elif mode == "holonomy":
                    # FIX #4: Use evolved alpha, not fixed
                    qgt_result = compute_qgt_holonomy_space(alpha, primes, engine)
                
                R = qgt_result["curvature_scalar"]
            except Exception as e:
                if verbose and step == 0:
                    print(f"  QGT computation failed: {e}")
                R = np.nan
        else:
            R = np.nan
        
        # Compute diagnostics
        S = float(compute_S_rep(psi))
        S_local = compute_local_S_rep(psi)
        grad_norm = compute_gradient_norm(S_local, laplacian)
        
        history.append({
            "step": step,
            "tau": tau,
            "R": R,
            "S_rep": S,
            "grad_norm": float(grad_norm),
        })
        
        # Evolve
        if step < n_steps:
            psi = gradient_flow_step(psi, H, dt)
            # FIX #5: Normalize after each step to keep ψ stable under L₆ evolution
            psi = psi / (jnp.linalg.norm(psi) + 1e-12)
    
    return pd.DataFrame(history)


def fit_parameters_to_state(
    psi_target: jnp.ndarray,
    generators: List[jnp.ndarray],
    base_psi: jnp.ndarray,
    n_iter: int = 20,
    lr: float = 0.05,
) -> jnp.ndarray:
    """
    Find θ such that ψ_target ≈ exp(Σ θᵃ Xₐ) ψ₀
    
    Uses gradient descent.
    """
    # Initialize at origin
    theta = jnp.zeros(len(generators))
    
    def loss_fn(th):
        psi_pred = build_exponential_state(th, generators, base_psi)
        return jnp.sum((psi_pred - psi_target)**2)
    
    # Gradient descent
    for _ in range(n_iter):
        grad = jax.grad(loss_fn)(theta)
        theta = theta - lr * grad
    
    return theta


def analyze_curvature_bursts(
    df: pd.DataFrame,
) -> Dict[str, Any]:
    """
    Analyze curvature spikes during evolution.
    """
    # Filter out NaN values
    df_valid = df[df["R"].notna()].copy()
    
    if len(df_valid) == 0:
        return {
            "max_R": 0.0,
            "max_R_abs": 0.0,
            "burst_time": 0.0,
            "n_samples": 0,
        }
    
    # Find maximum curvature
    df_valid["R_abs"] = df_valid["R"].abs()
    idx_max = df_valid["R_abs"].idxmax()
    
    return {
        "max_R": float(df_valid.loc[idx_max, "R"]),
        "max_R_abs": float(df_valid.loc[idx_max, "R_abs"]),
        "burst_time": float(df_valid.loc[idx_max, "tau"]),
        "n_samples": len(df_valid),
    }


# =============================================================================
# Multi-Plaquette Scanner with Next-Gen Parameterizations
# =============================================================================

def scan_all_plaquettes_nextgen(
    mode: str = "exponential",
    n_steps: int = 100,
    sample_every: int = 10,
    verbose: bool = True,
) -> pd.DataFrame:
    """
    Scan all 9 plaquettes with next-gen parameterization.
    
    Args:
        mode: "exponential" or "holonomy"
        n_steps: Evolution steps
        sample_every: QGT sampling frequency
        verbose: Print progress
    
    Returns:
        DataFrame with curvature burst analysis for all particles
    """
    engine = CliffordEngine(seed=42)
    results = []
    
    if verbose:
        print("=" * 80)
        print(f"Next-Gen Curvature Scanner (mode={mode})")
        print("=" * 80)
        print()
    
    for particle in ALL_PARTICLES:
        primes_closed = PLAQUETTES[particle]
        primes = primes_closed[:-1]  # Remove closing prime
        mass = PDG_MASSES[particle]
        
        if verbose:
            print(f"Scanning {particle:3s} (m={mass:10.2f} MeV): {primes}")
        
        # Track curvature during flow
        df_time = track_curvature_during_flow(
            primes, engine, mode, n_steps, sample_every=sample_every, verbose=False
        )
        
        # Analyze bursts
        bursts = analyze_curvature_bursts(df_time)
        
        results.append({
            "particle": particle,
            "mass_MeV": mass,
            "max_R": bursts["max_R"],
            "max_R_abs": bursts["max_R_abs"],
            "burst_time": bursts["burst_time"],
            "n_samples": bursts["n_samples"],
        })
        
        if verbose:
            print(f"  → max|R|={bursts['max_R_abs']:.2e} at τ={bursts['burst_time']:.3f}")
            print()
    
    df = pd.DataFrame(results)
    
    # Save results
    if verbose:
        log_dir = Path(__file__).parent / "logs"
        log_dir.mkdir(exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_path = log_dir / f"nextgen_curvature_{mode}_{timestamp}.csv"
        df.to_csv(csv_path, index=False)
        print(f"Results saved to: {csv_path}")
        print()
    
    return df


# =============================================================================
# Demo
# =============================================================================

def demo_nextgen_curvature():
    """Demo all four upgrades."""
    print("=" * 80)
    print("Next-Generation Curvature Parameterization")
    print("=" * 80)
    print()
    print("Testing four major upgrades:")
    print("  A. Exponential generators (non-linear)")
    print("  B. Holonomy-space parameterization (gauge field geometry)")
    print("  C. Spinor states (full Clifford DOF)")
    print("  D. Time-dependent QGT (transient curvature)")
    print()
    print("=" * 80)
    print()
    
    # Test Upgrade A: Exponential generators
    print("Testing Upgrade A: Exponential Generators")
    print("-" * 80)
    
    engine = CliffordEngine(seed=42)
    primes = [2, 3, 5, 7]
    generators = build_edge_generators(primes, engine)
    
    # FIX #7: Break symmetry in base state using prime logarithms
    base_psi = jnp.log(jnp.array(primes, dtype=float)) 
    base_psi = base_psi / jnp.linalg.norm(base_psi)
    
    theta_test = jnp.array([0.1, 0.2, 0.3, 0.4])
    
    result_exp = compute_qgt_exponential(theta_test, generators, base_psi, engine, primes)
    print(f"  R (exponential) = {result_exp['curvature_scalar']:.2e}")
    print(f"  Fisher trace = {jnp.trace(result_exp['fisher']):.2e}")
    print()
    
    # Test Upgrade B: Holonomy-space
    print("Testing Upgrade B: Holonomy-Space Parameterization")
    print("-" * 80)
    
    alpha_test = jnp.array([1.0, 1.0, 1.0, 1.0])
    result_hol = compute_qgt_holonomy_space(alpha_test, primes, engine)
    print(f"  R (holonomy) = {result_hol['curvature_scalar']:.2e}")
    print(f"  Fisher trace = {jnp.trace(result_hol['fisher']):.2e}")
    print()
    
    # Scan all plaquettes with exponential mode
    print("=" * 80)
    print("Multi-Plaquette Scan: Exponential Mode")
    print("=" * 80)
    print()
    
    df_exp = scan_all_plaquettes_nextgen(mode="exponential", n_steps=50, sample_every=10)
    
    print()
    print("Results Summary:")
    print(df_exp.sort_values("mass_MeV")[["particle", "mass_MeV", "max_R_abs", "burst_time"]])
    print()
    
    print("=" * 80)
    print("Demo Complete!")
    print("=" * 80)


if __name__ == "__main__":
    demo_nextgen_curvature()
