# OID_analysis/utils/analysis.py
"""
CPU-only analysis utilities: POD, OID, PCD, zone statistics, comparison.
No pycuda dependency -- import from any environment.
"""
from __future__ import annotations

import numpy as np
from typing import Dict, List, Optional, Tuple


# ---------------------------------------------------------------------------
# POD (Proper Orthogonal Decomposition)
# ---------------------------------------------------------------------------

def compute_pod(
    snapshots: np.ndarray,
    rank: Optional[int] = None,
) -> Dict:
    """Compute POD of snapshot matrix.

    Args:
        snapshots: (N, DOF) array, N snapshots of DOF-dimensional field.
        rank: Truncation rank. If None, keep all.

    Returns:
        dict with keys:
            modes: (DOF, r) spatial modes
            coefs: (N, r) time coefficients
            S: (r,) singular values
            energy: (r,) relative energy fraction per mode
            cum_energy: (r,) cumulative energy
            mean: (DOF,) mean field
    """
    mean = np.mean(snapshots, axis=0)
    Q = snapshots - mean  # (N, DOF)
    
    # Prefer method-of-snapshots if N < DOF
    N, DOF = Q.shape
    if N < DOF:
        # Method of snapshots: SVD on (N x N) covariance
        C = Q @ Q.T  # (N, N)
        eigvals, eigvecs = np.linalg.eigh(C)
        idx = np.argsort(eigvals)[::-1]
        eigvals = eigvals[idx]
        eigvecs = eigvecs[:, idx]
        
        # Recover modes
        S = np.sqrt(np.maximum(eigvals, 0))
        V = eigvecs
        # phi_i = (1/sigma_i) * Q^T * v_i
        modes = (Q.T @ V) / (S + 1e-30)  # (DOF, N)
        coefs = V * S  # (N, N)
    else:
        # Direct SVD
        U, S, Vt = np.linalg.svd(Q, full_matrices=False)
        modes = Vt.T  # (DOF, N)
        coefs = U * S  # (N, N)
    
    # Energy is always normalized by the complete snapshot spectrum.  The old
    # implementation normalized after truncation, making every retained rank
    # incorrectly appear to capture 100% of fluctuation energy.
    total_energy = np.sum(S ** 2)

    # Truncate only after preserving the full-spectrum denominator.
    max_rank = min(N, DOF)
    if rank is not None:
        if type(rank) is not int or rank < 1 or rank > max_rank:
            raise ValueError(f"rank must be in [1, {max_rank}]")
        modes = modes[:, :rank]
        S = S[:rank]
        coefs = coefs[:, :rank]

    energy = (S ** 2) / total_energy if total_energy > 0 else np.zeros_like(S)
    cum_energy = np.cumsum(energy)
    
    return {
        "modes": modes,
        "coefs": coefs,
        "S": S,
        "energy": energy,
        "cum_energy": cum_energy,
        "mean": mean,
    }


# ---------------------------------------------------------------------------
# OID (Observable-Inferred Decomposition) -- cross-covariance SVD
# ---------------------------------------------------------------------------

def compute_force_oid(
    pod_coefs: np.ndarray,
    y_force: np.ndarray,
) -> Dict:
    """Force-OID: cross-covariance SVD between POD coefs and force observable.

    Args:
        pod_coefs: (N, r) standardized correction-field POD coefficients.
        y_force: (N, m) standardized force observable.

    Returns:
        dict with keys:
            U: (r, r) OID rotation matrix
            S: (r,) singular values of cross-covariance
            Vt: (m, m) right singular vectors
            z: (N, r) OID coordinates
            cum_corr: (r,) cumulative correlation fraction
    """
    N = pod_coefs.shape[0]
    C_AY = (1.0 / N) * pod_coefs.T @ y_force  # (r, m)
    U, S, Vt = np.linalg.svd(C_AY, full_matrices=False)
    z = pod_coefs @ U  # (N, r)
    
    total_corr = np.sum(S)
    cum_corr = np.cumsum(S) / total_corr if total_corr > 0 else np.ones_like(S)
    
    return {
        "U": U,
        "S": S,
        "Vt": Vt,
        "z": z,
        "cum_corr": cum_corr,
    }


def compute_signature_oid(
    pod_coefs: np.ndarray,
    y_sig: np.ndarray,
) -> Dict:
    """Signature-OID: cross-covariance SVD between POD coefs and delayed sensor error.

    Identical to force-OID but with different observable. Kept as separate
    function for documentation clarity.
    """
    return compute_force_oid(pod_coefs, y_sig)


# ---------------------------------------------------------------------------
# PCD (Pattern-Constrained Decomposition) -- whitened cross-correlation
# ---------------------------------------------------------------------------

def compute_pcd(
    pod_coefs: np.ndarray,
    p_sig: np.ndarray,
    tikhonov_eps: float = 1e-6,
) -> Dict:
    """PCD-style whitened cross-correlation decomposition.

    Implements Lyu23-style canonical correlation in POD subspace.

    Args:
        pod_coefs: (N, r) standardized correction-field POD coefficients.
        p_sig: (N, m') delayed signature observable (e.g. p_sig_full).
        tikhonov_eps: Tikhonov regularization for near-singular matrices.

    Returns:
        dict with keys:
            W: (r, r) PCD weight matrix (mapping POD coefs -> PCD coords)
            z_pcd: (N, r) PCD coordinates
            S: (r,) singular values of whitened cross-correlation
            cum_corr: (r,) cumulative correlation fraction
    """
    import scipy.linalg
    N = pod_coefs.shape[0]
    r = pod_coefs.shape[1]
    
    # Covariance matrices
    C_AA = (1.0 / N) * pod_coefs.T @ pod_coefs  # (r, r)
    C_PP = (1.0 / N) * p_sig.T @ p_sig  # (m', m')
    C_AP = (1.0 / N) * pod_coefs.T @ p_sig  # (r, m')
    
    # Whitening with Tikhonov
    C_AA_reg = C_AA + tikhonov_eps * np.eye(r)
    C_PP_reg = C_PP + tikhonov_eps * np.eye(p_sig.shape[1])
    
    C_AA_inv_half = scipy.linalg.sqrtm(np.linalg.inv(C_AA_reg)).real
    C_PP_inv_half = scipy.linalg.sqrtm(np.linalg.inv(C_PP_reg)).real
    
    # Whitened cross-correlation
    K = C_AA_inv_half @ C_AP @ C_PP_inv_half
    U, S, Vt = np.linalg.svd(K, full_matrices=False)
    
    W = C_AA_inv_half @ U  # (r, r)
    z_pcd = pod_coefs @ W  # (N, r)
    
    total_corr = np.sum(S)
    cum_corr = np.cumsum(S) / total_corr if total_corr > 0 else np.ones_like(S)
    
    return {
        "W": W,
        "z_pcd": z_pcd,
        "S": S,
        "U_raw": U,
        "cum_corr": cum_corr,
    }


# ---------------------------------------------------------------------------
# Zone statistics
# ---------------------------------------------------------------------------

def compute_zone_statistics(
    u: np.ndarray,
    v: np.ndarray,
    zone_mask: np.ndarray,
) -> Dict:
    """Compute coarse-grained statistics within a zone.

    Args:
        u: (N, ny, nx) x-velocity time series
        v: (N, ny, nx) y-velocity time series
        zone_mask: (ny, nx) bool mask for the zone

    Returns:
        dict with keys:
            kinetic_energy_mean: scalar (time-mean within zone)
            enstrophy_mean: scalar
            upper_lower_asymmetry: scalar
            centerline_shift: scalar
    """
    N = u.shape[0]
    mask = zone_mask.astype(bool)
    area = np.sum(mask)
    
    if area == 0:
        return {"area": 0, "warning": "empty zone"}
    
    # Kinetic energy per snapshot
    ke_per_step = 0.5 * np.sum((u[:, mask] ** 2 + v[:, mask] ** 2), axis=1) / area
    ke_mean = float(np.mean(ke_per_step))
    
    # Vorticity via finite difference (central, on grid)
    omega_ts = np.zeros((N, *u.shape[1:]), dtype=np.float64)
    for t in range(N):
        u_t = u[t].astype(np.float64)
        v_t = v[t].astype(np.float64)
        omega_ts[t] = np.gradient(v_t, axis=1) - np.gradient(u_t, axis=0)
    
    enstrophy_per_step = 0.5 * np.sum(omega_ts[:, mask] ** 2, axis=1) / area
    enstrophy_mean = float(np.mean(enstrophy_per_step))
    
    # Upper-lower asymmetry (split along center y)
    ny = u.shape[1]
    cy = (ny - 1) / 2.0
    upper_mask = mask.copy()
    upper_mask[:int(cy), :] = False
    lower_mask = mask.copy()
    lower_mask[int(cy):, :] = False
    
    upper_omega = np.mean(omega_ts[:, upper_mask]) if np.sum(upper_mask) > 0 else 0.0
    lower_omega = np.mean(omega_ts[:, lower_mask]) if np.sum(lower_mask) > 0 else 0.0
    total = max(abs(upper_omega) + abs(lower_omega), 1e-30)
    asymmetry = float((upper_omega - lower_omega) / total)
    
    # Centerline shift: find y where mean u(y) is max on the centerline
    mean_u = np.mean(u, axis=0)  # (ny, nx)
    center_x_idx = mean_u.shape[1] // 2
    u_profile = mean_u[:, center_x_idx]
    shift_y = float(np.argmax(u_profile) - cy)
    
    return {
        "area": int(area),
        "kinetic_energy_mean": ke_mean,
        "enstrophy_mean": enstrophy_mean,
        "upper_lower_asymmetry": asymmetry,
        "centerline_shift_lattice": shift_y,
    }


def compute_recirculation_metrics(
    mean_u: np.ndarray,
    centerline_y: int,
) -> Dict:
    """Compute recirculation zone metrics from mean x-velocity field.

    Args:
        mean_u: (ny, nx) time-mean x-velocity
        centerline_y: y-index of the centerline

    Returns:
        dict with keys:
            Lr: recirculation length (max x where u=0 on centerline)
            Ar: recirculation area (pixels where u<0)
    """
    ny, nx = mean_u.shape
    # Recirculation length: furthest downstream x on centerline where mean_u < 0
    u_cl = mean_u[centerline_y, :]
    neg_indices = np.where(u_cl < 0)[0]
    Lr = float(neg_indices[-1]) if len(neg_indices) > 0 else 0.0
    
    # Recirculation area: pixels where mean_u < 0
    Ar = float(np.sum(mean_u < 0))
    
    # Centerline shift: y-offset of max u in near wake
    wake_region = mean_u[:, nx//4:nx//2]
    max_uy = np.unravel_index(np.argmax(wake_region), wake_region.shape)
    cl_shift = float(max_uy[0] - centerline_y)
    
    return {
        "Lr_lattice": Lr,
        "Ar_pixels": Ar,
        "centerline_shift_lattice": cl_shift,
    }


# ---------------------------------------------------------------------------
# Standardization
# ---------------------------------------------------------------------------

def standardize(X: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Z-score standardization along time axis (axis=0)."""
    mean = np.mean(X, axis=0, keepdims=True)
    std = np.std(X, axis=0, keepdims=True)
    std = np.where(std < 1e-12, 1.0, std)
    X_std = (X - mean) / std
    return X_std, mean, std


def reconstruct_oid_modes(
    pod_modes: np.ndarray,
    U_oid: np.ndarray,
    coefficient_scale: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Reconstruct physical PLS-type directions from standardized POD space.

    If decomposition directions ``U_oid`` were fitted to standardized POD
    coefficients ``(A - mean) / scale``, the physical field direction is
    ``Phi @ diag(scale) @ U``.  ``coefficient_scale=None`` is retained only for
    legacy callers whose coefficients were not standardized.
    """
    modes = np.asarray(pod_modes)
    directions = np.asarray(U_oid)
    if coefficient_scale is None:
        return modes @ directions
    scale = np.asarray(coefficient_scale).reshape(-1)
    if modes.shape[1] != len(scale) or directions.shape[0] != len(scale):
        raise ValueError("pod_modes, coefficient_scale, and U_oid dimensions differ")
    return modes @ (scale[:, None] * directions)


# ---------------------------------------------------------------------------
# Comparison helpers
# ---------------------------------------------------------------------------

def compare_predictions(
    y_true: np.ndarray,
    y_pred_pod: np.ndarray,
    y_pred_oid: np.ndarray,
    y_pred_pcd: Optional[np.ndarray] = None,
) -> Dict:
    """Compare prediction R^2 between POD, OID, and optionally PCD."""
    from sklearn.linear_model import LinearRegression
    
    def r2_score(y_t, y_p):
        ss_res = np.sum((y_t - y_p) ** 2)
        ss_tot = np.sum((y_t - np.mean(y_t, axis=0, keepdims=True)) ** 2)
        return 1.0 - ss_res / (ss_tot + 1e-30)
    
    result = {
        "pod_r2": r2_score(y_true, y_pred_pod),
        "oid_r2": r2_score(y_true, y_pred_oid),
    }
    if y_pred_pcd is not None:
        result["pcd_r2"] = r2_score(y_true, y_pred_pcd)
    
    return result
