# OID_analysis/tests/test_analysis.py
"""Unit tests for utils/analysis.py — POD, OID, PCD, and helpers.

References: Sch12 = Schlegel et al. (2012) JFM.
"""
import numpy as np
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
from OID_analysis.utils.analysis import (
    compute_pod, compute_force_oid, compute_pcd,
    standardize, reconstruct_oid_modes,
)

# ---------------------------------------------------------------------------
# Test 1: POD energy conservation (Sch12 Eq 2.3)
# ---------------------------------------------------------------------------

def test_pod_energy_conservation():
    """POD eigenvalues sum to total fluctuation energy (Sch12 Eq 2.3).

    For method-of-snapshots: sum(S^2) = trace(Q^T Q)/N = total variance.
    We check that POD energy fractions sum to 1 and the total energy
    captured matches a direct full-field computation via full-rank POD.
    """
    np.random.seed(42)
    N, DOF = 200, 500  # use DOF > N to trigger method-of-snapshots
    t = np.linspace(0, 4*np.pi, N)
    mode1 = np.sin(t).reshape(-1, 1) @ np.random.randn(1, DOF) * 3.0
    mode2 = np.cos(2*t).reshape(-1, 1) @ np.random.randn(1, DOF) * 2.0
    mean = np.random.randn(1, DOF)
    snapshots = mean + mode1 + mode2 + np.random.randn(N, DOF) * 0.05

    # Full POD (no truncation)
    pod = compute_pod(snapshots, rank=None)
    assert abs(np.sum(pod["energy"]) - 1.0) < 1e-10, \
        f"Energy fractions don't sum to 1: {np.sum(pod['energy'])}"

    # Cumulative energy should reach 1.0 at full rank
    assert pod["cum_energy"][-1] > 0.9999

    # First 3 modes should capture >95% (2 true modes + mean removal)
    assert pod["cum_energy"][2] > 0.95, \
        f"Energy capture too low: {pod['cum_energy'][2]:.4f}"

# ---------------------------------------------------------------------------
# Test 2: OID cross-covariance (Sch12 Eq 2.8, 2.21)
# ---------------------------------------------------------------------------

def test_oid_cross_covariance():
    """OID finds direction that best correlates with observable.

    Sch12 Eq 2.8: b = C a. Our C_AY = (1/N) A^T Y estimates C.
    U columns from SVD give directions in A-space that maximize correlation.
    """
    np.random.seed(123)
    N, r, m = 500, 10, 3  # r > m case is the common scenario
    A = np.random.randn(N, r)
    # Y depends primarily on A[:, 3] + some from A[:, 1]
    Y = A[:, 3:4] * 2.0 + A[:, 1:2] * 0.5 + np.random.randn(N, 1) * 0.01
    # Extend Y to m columns for richness
    Y = np.hstack([Y, np.random.randn(N, m-1) * 0.1])
    A_std, _, _ = standardize(A)
    Y_std, _, _ = standardize(Y)

    oid = compute_force_oid(A_std, Y_std)

    # First OID coordinate should strongly correlate with Y[:,0]
    z1 = oid["z"][:, 0]
    corr_z1_y = np.corrcoef(z1, Y_std[:, 0])[0, 1]
    assert abs(corr_z1_y) > 0.7, f"OID z1-Y correlation too low: {corr_z1_y:.3f}"

    # S values should be sorted descending
    assert np.all(np.diff(oid["S"]) <= 0), "S values not sorted descending"

    # U^T @ U = I_m (since U is r×m, semiorthogonal when r > m)
    U = oid["U"]  # shape (r, min(r,m)) = (10, 3)
    assert U.shape[1] == m, f"U shape wrong: {U.shape}"
    assert np.allclose(U.T @ U, np.eye(m), atol=1e-10), \
        f"U^T @ U != I_m: max error {np.max(np.abs(U.T @ U - np.eye(m)))}"

# ---------------------------------------------------------------------------
# Test 3: OID coordinate reconstruction via POD modes
# ---------------------------------------------------------------------------

def test_oid_field_reconstruction():
    """OID coordinates reconstruct the field through POD and OID modes.

    z = A @ U, psi_OID = Phi @ U, so A @ Phi^T = z @ U^T @ Phi^T.
    If U is square (r==m), then A @ Phi^T = z @ psi_OID^T (exact).
    If r > m, reconstruction is best m-dimensional approximation.
    """
    np.random.seed(42)
    N, DOF, r = 100, 200, 6
    snapshots = np.random.randn(N, DOF)
    pod = compute_pod(snapshots, rank=r)
    A = pod["coefs"]  # (N, r)
    Phi = pod["modes"]  # (DOF, r)
    mean = pod["mean"]  # (DOF,)

    # Y uses all r directions to make U square
    Y = A @ np.random.randn(r, r) + np.random.randn(N, r) * 0.01
    A_std, _, _ = standardize(A)
    Y_std, _, _ = standardize(Y)
    oid = compute_force_oid(A_std, Y_std)
    U = oid["U"]  # (r, r) since Y has r columns
    z = oid["z"]  # (N, r)

    # Reconstruct in standardized space: z @ psi_OID.T = A_std @ U @ U^T @ Phi.T
    # Since U is square orthogonal, this = A_std @ Phi.T
    # Standardized POD coordinates map to physical directions through Phi D U.
    _, _, a_scale = standardize(A)
    psi_oid = reconstruct_oid_modes(Phi, U, a_scale)  # (DOF, r)
    q_std_oid = z @ psi_oid.T
    q_std_direct = A_std @ np.diag(a_scale.reshape(-1)) @ Phi.T
    rel_err = np.mean((q_std_oid - q_std_direct)**2) / \
              (np.mean(q_std_direct**2) + 1e-30)
    assert rel_err < 1e-6, f"Reconstruction error: {rel_err:.10f}"

# ---------------------------------------------------------------------------
# Test 4: PCD whitening correlation
# ---------------------------------------------------------------------------

def test_pcd_whitening():
    """PCD whitened CCA finds canonical correlations (Sch12 §2.4 concept)."""
    np.random.seed(99)
    N, r, m = 300, 6, 3
    A = np.random.randn(N, r)
    Y = A @ np.random.randn(r, m) * 0.5 + np.random.randn(N, m) * 0.5
    A_std, _, _ = standardize(A)
    Y_std, _, _ = standardize(Y)

    pcd = compute_pcd(A_std, Y_std, tikhonov_eps=1e-4)

    # Cumulative correlation should be finite
    assert pcd["cum_corr"][-1] > 0 and not np.isnan(pcd["cum_corr"][-1]), \
        f"PCD total correlation NaN or zero: {pcd['cum_corr'][-1]}"

    # First PCD coordinate should have non-zero variance
    z_first = pcd["z_pcd"][:, 0]
    assert np.std(z_first) > 0.01, f"First PCD coordinate near-zero std: {np.std(z_first)}"

    # S values sorted descending
    assert np.all(np.diff(pcd["S"]) <= 0), "PCD S values not sorted"

# ---------------------------------------------------------------------------
# Test 5: Standardize zero-variance channels
# ---------------------------------------------------------------------------

def test_standardize_zero_std():
    """Channels with zero variance: no division by zero (clamp at 1e-12)."""
    X = np.array([[1.0, 0.5, 0.0],
                   [1.0, 0.5, 0.0],
                   [1.0, 0.5, 0.0]], dtype=np.float64)
    X_std, mean, std = standardize(X)
    assert np.all(np.isfinite(X_std)), "NaN in standardized array"
    assert np.allclose(X_std[:, 2], 0.0), "Zero-var channel must be zero after centering"
    assert np.allclose(X_std[:, 0], 0.0), "Constant channel must be zero"

# ---------------------------------------------------------------------------
# Test 6: POD snapshot method (N < DOF case)
# ---------------------------------------------------------------------------

def test_pod_method_of_snapshots():
    """When N < DOF, method-of-snapshots via (N,N) eigh is used and correct."""
    np.random.seed(7)
    N, DOF = 50, 1000
    r_true = 10
    U_true = np.linalg.qr(np.random.randn(DOF, r_true))[0]
    S_true = np.exp(-np.arange(r_true) / 3.0)
    coefs_true = np.random.randn(N, r_true) * S_true[np.newaxis, :]
    snapshots = coefs_true @ U_true.T + np.random.randn(N, DOF) * 0.01

    pod = compute_pod(snapshots, rank=15)
    assert pod["cum_energy"][r_true - 1] > 0.95, \
        f"Energy capture too low: {pod['cum_energy'][r_true-1]:.4f}"

# ---------------------------------------------------------------------------
# Test 7: OID cum_corr monotonic
# ---------------------------------------------------------------------------

def test_oid_cum_corr_monotonic():
    """cum_corr should monotonically increase and reach 1.0 on the last mode."""
    np.random.seed(456)
    N, r, m = 200, 12, 5
    A = np.random.randn(N, r)
    Y = A[:, :3] @ np.random.randn(3, m) + np.random.randn(N, m) * 0.3
    A_std, _, _ = standardize(A)
    Y_std, _, _ = standardize(Y)

    oid = compute_force_oid(A_std, Y_std)
    cum = oid["cum_corr"]
    assert np.all(np.diff(cum) >= 0), "cum_corr not monotonic"
    assert abs(cum[-1] - 1.0) < 1e-10, f"Last cum_corr != 1: {cum[-1]}"

# ---------------------------------------------------------------------------
# Run all tests
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    tests = [
        ("POD energy conservation", test_pod_energy_conservation),
        ("OID cross-covariance", test_oid_cross_covariance),
        ("OID field reconstruction", test_oid_field_reconstruction),
        ("PCD whitening", test_pcd_whitening),
        ("Standardize zero-std", test_standardize_zero_std),
        ("POD method of snapshots", test_pod_method_of_snapshots),
        ("OID cum_corr monotonic", test_oid_cum_corr_monotonic),
    ]
    passed = 0
    for name, fn in tests:
        try:
            fn()
            print(f"  PASS: {name}")
            passed += 1
        except AssertionError as e:
            print(f"  FAIL: {name} — {e}")
        except Exception as e:
            print(f"  ERROR: {name} — {type(e).__name__}: {e}")
    print(f"\n{passed}/{len(tests)} tests passed")
