"""Mathematically explicit correction-field POD and observable decompositions."""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import scipy.linalg


@dataclass(frozen=True)
class Standardizer:
    mean: np.ndarray
    scale: np.ndarray

    @classmethod
    def fit(cls, values: np.ndarray, eps: float = 1e-12) -> "Standardizer":
        array = np.asarray(values, dtype=float)
        scale = array.std(axis=0)
        return cls(array.mean(axis=0), np.where(scale < eps, 1.0, scale))

    def transform(self, values: np.ndarray) -> np.ndarray:
        return (np.asarray(values) - self.mean) / self.scale


@dataclass(frozen=True)
class PODModel:
    mean: np.ndarray
    modes: np.ndarray
    singular_values: np.ndarray
    coefficients: np.ndarray
    energy: np.ndarray
    cumulative_energy: np.ndarray
    total_energy: float
    metric: np.ndarray

    def transform(self, snapshots: np.ndarray) -> np.ndarray:
        centered = np.asarray(snapshots, dtype=float) - self.mean
        if self.metric.ndim == 1:
            return (centered * self.metric) @ self.modes
        return centered @ self.metric @ self.modes

    def inverse_transform(self, coefficients: np.ndarray) -> np.ndarray:
        return self.mean + np.asarray(coefficients, dtype=float) @ self.modes.T


def _as_metric(metric: np.ndarray | None, dimension: int, *, diagonal: bool = False) -> np.ndarray:
    if metric is None:
        return np.ones(dimension) if diagonal else np.eye(dimension)
    value = np.asarray(metric, dtype=float)
    if value.ndim == 1:
        if len(value) != dimension or np.any(value < 0) or not np.isfinite(value).all():
            raise ValueError("diagonal metric must contain finite nonnegative weights")
        if diagonal and not np.any(value > 0):
            raise ValueError("metric has no positive-weight degrees of freedom")
        return value
    if value.shape != (dimension, dimension):
        raise ValueError("metric must be a symmetric square matrix")
    _guard_dense_metric(value)
    if not np.allclose(value, value.T):
        raise ValueError("metric must be a symmetric square matrix")
    return value


# Dense metrics are useful for low-dimensional observables, but their eigensolver
# workspace and roots are quadratic in the field dimension. Large fields must
# use vector weights; this guard is checked before making another dense copy.
_DENSE_METRIC_MAX_DIMENSION = 4096
_DENSE_METRIC_MAX_BYTES = 256 * 1024**2


def _guard_dense_metric(metric: np.ndarray) -> None:
    if metric.shape[0] > _DENSE_METRIC_MAX_DIMENSION or metric.nbytes > _DENSE_METRIC_MAX_BYTES:
        raise ValueError(
            "dense metric is too large; pass diagonal weights as a one-dimensional vector"
        )


def _metric_root(metric: np.ndarray, *, positive_definite: bool = False) -> tuple[np.ndarray, np.ndarray]:
    if metric.ndim == 1:
        tolerance = np.finfo(float).eps * max(len(metric), 1) * max(float(np.max(metric)), 1.0)
        if positive_definite and np.any(metric <= tolerance):
            raise ValueError("metric must be positive definite")
        positive = metric > tolerance
        root = np.zeros_like(metric)
        inverse_root = np.zeros_like(metric)
        root[positive] = np.sqrt(metric[positive])
        inverse_root[positive] = 1.0 / root[positive]
        return root, inverse_root
    _guard_dense_metric(metric)
    eigenvalues, eigenvectors = scipy.linalg.eigh(metric, check_finite=True)
    tolerance = np.finfo(float).eps * max(metric.shape) * max(float(np.max(np.abs(eigenvalues))), 1.0)
    if eigenvalues[0] < -tolerance or (positive_definite and eigenvalues[0] <= tolerance):
        kind = "positive definite" if positive_definite else "positive semidefinite"
        raise ValueError(f"metric must be {kind}")
    positive = eigenvalues > tolerance
    root = (eigenvectors[:, positive] * np.sqrt(eigenvalues[positive])) @ eigenvectors[:, positive].T
    inverse_root = (eigenvectors[:, positive] * eigenvalues[positive]**-0.5) @ eigenvectors[:, positive].T
    return root, inverse_root


def _fix_mode_signs(modes: np.ndarray, coefficients: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Choose a repeatable sign without changing a POD factorization."""
    if modes.shape[1] == 0:
        return modes, coefficients
    pivots = np.argmax(np.abs(modes), axis=0)
    signs = np.sign(modes[pivots, np.arange(modes.shape[1])])
    signs[signs == 0] = 1.0
    return modes * signs, coefficients * signs


def fit_pod(snapshots: np.ndarray, rank: int, weights: np.ndarray | None = None) -> PODModel:
    """Fit field POD by the deterministic weighted method of snapshots.

    For ``N`` snapshots with ``N << D``, this diagonal-metric path forms only
    ``K = (X sqrt(w)) (X sqrt(w))^T`` (``N x N``), never a ``D x D`` metric.
    Zero-weight coordinates are explicitly zero in every returned mode.
    """
    data = np.asarray(snapshots, dtype=float)
    if data.ndim != 2 or data.shape[0] < 2 or rank < 1:
        raise ValueError("invalid snapshot matrix or rank")
    metric = _as_metric(weights, data.shape[1], diagonal=True)
    mean = data.mean(axis=0)
    centered = data - mean

    if metric.ndim == 1:
        weighted = centered * np.sqrt(metric)
        gram = weighted @ weighted.T
        eigenvalues, eigenvectors = scipy.linalg.eigh(gram, check_finite=False)
        order = np.argsort(eigenvalues)[::-1]
        eigenvalues = np.maximum(eigenvalues[order], 0.0)
        eigenvectors = eigenvectors[:, order]
        full_s = np.sqrt(eigenvalues)
        # A Gram eigensolve resolves squared singular values; threshold in that
        # spectrum so centering/null directions are not promoted by sqrt noise.
        eigenvalue_tolerance = (np.finfo(float).eps * max(centered.shape) *
                                (eigenvalues[0] if len(eigenvalues) else 1.0))
        numerical_rank = int(np.sum(eigenvalues > eigenvalue_tolerance))
        retained_rank = min(int(rank), numerical_rank)
        if retained_rank < 1:
            raise ValueError("weighted snapshot matrix has zero numerical rank")
        retained = full_s[:retained_rank]
        modes = np.zeros((data.shape[1], retained_rank), dtype=float)
        positive = metric > 0
        modes[positive] = centered[:, positive].T @ eigenvectors[:, :retained_rank] / retained
        coefficients = eigenvectors[:, :retained_rank] * retained
        modes, coefficients = _fix_mode_signs(modes, coefficients)
        total = float(np.trace(gram))
    else:
        _guard_dense_metric(metric)
        root, inverse_root = _metric_root(metric)
        _, full_s, full_vt = np.linalg.svd(centered @ root, full_matrices=False)
        tolerance = np.finfo(float).eps * max(centered.shape) * (full_s[0] if len(full_s) else 1.0)
        numerical_rank = int(np.sum(full_s > tolerance))
        retained_rank = min(int(rank), numerical_rank)
        if retained_rank < 1:
            raise ValueError("weighted snapshot matrix has zero numerical rank")
        retained = full_s[:retained_rank]
        modes = inverse_root @ full_vt[:retained_rank].T
        coefficients = centered @ metric @ modes
        modes, coefficients = _fix_mode_signs(modes, coefficients)
        total = float(np.sum(full_s**2))
    energy = retained**2 / total if total > 0 else np.zeros(retained_rank)
    return PODModel(mean, modes, retained, coefficients, energy, np.cumsum(energy), total, metric)


def fit_observable_pod(observables: np.ndarray, rank: int, metric: np.ndarray | None = None) -> PODModel:
    """POD for observables under a diagonal or guarded dense PSD metric."""
    data = np.asarray(observables, dtype=float)
    if data.ndim != 2 or data.shape[0] < 2 or rank < 1:
        raise ValueError("invalid observable matrix or rank")
    g = _as_metric(metric, data.shape[1])
    root, inverse_root = _metric_root(g)
    mean = data.mean(axis=0)
    centered = data - mean
    weighted = centered * root if g.ndim == 1 else centered @ root
    _, full_s, full_vt = np.linalg.svd(weighted, full_matrices=False)
    tolerance = np.finfo(float).eps * max(centered.shape) * (full_s[0] if len(full_s) else 1.0)
    numerical_rank = int(np.sum(full_s > tolerance))
    retained_rank = min(int(rank), numerical_rank)
    if retained_rank < 1:
        raise ValueError("observable matrix has zero metric numerical rank")
    retained = full_s[:retained_rank]
    if g.ndim == 1:
        modes = inverse_root[:, None] * full_vt[:retained_rank].T
        coefficients = (centered * g) @ modes
    else:
        modes = inverse_root @ full_vt[:retained_rank].T
        coefficients = centered @ g @ modes
    modes, coefficients = _fix_mode_signs(modes, coefficients)
    total = float(np.sum(full_s**2))
    energy = retained**2 / total if total > 0 else np.zeros(retained_rank)
    return PODModel(mean, modes, retained, coefficients, energy, np.cumsum(energy), total, g)


@dataclass(frozen=True)
class OIDInverse:
    """Finite-dimensional Schlegel inverse and its algebraic diagnostics."""

    propagator: np.ndarray
    inverse: np.ndarray
    projection: np.ndarray
    observable_basis: np.ndarray
    singular_values: np.ndarray
    rank: int
    rank_tolerance: float
    identity_defect: float
    projection_defect: float
    propagation_defect: float
    residual_covariance: np.ndarray
    residual_observable_covariance: np.ndarray


@dataclass(frozen=True)
class SchlegelOID:
    observable_pod: PODModel
    propagator: np.ndarray
    lr: OIDInverse
    le: OIDInverse
    lr_modes: np.ndarray
    le_modes: np.ndarray
    state_mean: np.ndarray
    rcond: float
    requested_observable_rank: int
    effective_observable_rank: int
    state_covariance_rank: int
    state_covariance_threshold: float
    state_covariance_smallest_retained_singular: float

    def observable_coordinates(self, observables: np.ndarray) -> np.ndarray:
        """Project measured observables onto the generated range coordinates."""
        return self.observable_pod.transform(observables) @ self.lr.observable_basis

    def predict_observables(self, state_coefficients: np.ndarray) -> np.ndarray:
        """Predict physical observables, including generated-range truncation."""
        centered = np.asarray(state_coefficients, dtype=float) - self.state_mean
        generated = centered @ self.lr.propagator.T
        full_coordinates = generated @ self.lr.observable_basis.T
        return self.observable_pod.inverse_transform(full_coordinates)

    def predict_state_coefficients(
        self, observables: np.ndarray, *, inverse: str = "lr"
    ) -> np.ndarray:
        """Predict centered state coefficients from generated observable coordinates."""
        model = self.lr if inverse == "lr" else self.le if inverse == "le" else None
        if model is None:
            raise ValueError("inverse must be 'lr' or 'le'")
        return self.state_mean + self.observable_coordinates(observables) @ model.inverse.T


def _numerical_rank(singular_values: np.ndarray, shape: tuple[int, int], rcond: float | None) -> tuple[int, float]:
    largest = float(singular_values[0]) if len(singular_values) else 0.0
    relative = np.finfo(float).eps * max(shape) if rcond is None else float(rcond)
    if relative < 0:
        raise ValueError("rcond must be nonnegative")
    tolerance = relative * largest
    # rcond=0 intentionally retains every strictly positive singular value.
    return int(np.sum(singular_values > tolerance)), tolerance


def _pinv(matrix: np.ndarray, rcond: float | None) -> np.ndarray:
    u, singular_values, vt = np.linalg.svd(matrix, full_matrices=False)
    rank, _ = _numerical_rank(singular_values, matrix.shape, rcond)
    if rank == 0:
        return np.zeros((matrix.shape[1], matrix.shape[0]), dtype=float)
    return (vt[:rank].T / singular_values[:rank]) @ u[:, :rank].T


def fit_schlegel_oid(
    state_coefficients: np.ndarray,
    observables: np.ndarray,
    *,
    observable_metric: np.ndarray | None = None,
    observable_rank: int | None = None,
    state_metric: np.ndarray | None = None,
    rcond: float | None = None,
    truncate: bool = False,
    state_modes: np.ndarray | None = None,
) -> SchlegelOID:
    """Construct exact LR and LE inverses on the generated observable range.

    Samples are rows.  After metric observable POD, ``b = C a`` uses
    ``C = Sigma_ba Sigma_aa^+``--never the raw cross-covariance.  LR is the
    covariance-weighted minimum mean-square inverse; LE is the minimum
    ``state_metric``-energy inverse.  If ``C`` is not onto the requested
    observable coordinates, the function fails closed unless ``truncate`` is
    true, in which case all identities are explicitly restricted to its SVD
    generated range.
    """
    a = np.asarray(state_coefficients, dtype=float)
    y = np.asarray(observables, dtype=float)
    if a.ndim != 2 or y.ndim != 2 or len(a) != len(y) or len(a) < 2:
        raise ValueError("state and observable samples must be aligned matrices")
    requested_observable_rank = min(y.shape) if observable_rank is None else int(observable_rank)
    relative_rcond = np.finfo(float).eps * max(a.shape) if rcond is None else float(rcond)
    if not np.isfinite(relative_rcond) or relative_rcond < 0:
        raise ValueError("rcond must be a finite nonnegative scalar")
    observable_pod = fit_observable_pod(y, requested_observable_rank, observable_metric)
    # ``fit_observable_pod`` supports rank-max requests by slicing to numerical
    # rank; all generated-range algebra uses that effective coordinate count.
    requested = observable_pod.modes.shape[1]
    ac = a - a.mean(axis=0)
    b = observable_pod.coefficients
    n = len(a)
    sigma_aa = ac.T @ ac / n
    sigma_ba = b.T @ ac / n
    covariance_singular_values = np.linalg.svd(sigma_aa, compute_uv=False)
    state_covariance_rank, state_covariance_threshold = _numerical_rank(
        covariance_singular_values, sigma_aa.shape, relative_rcond
    )
    if state_covariance_rank < 1:
        raise np.linalg.LinAlgError("state covariance has zero numerical rank")
    propagator = sigma_ba @ _pinv(sigma_aa, relative_rcond)

    u, singular_values, vt = np.linalg.svd(propagator, full_matrices=False)
    generated_rank, rank_tolerance = _numerical_rank(
        singular_values, propagator.shape, relative_rcond
    )
    if generated_rank < requested and not truncate:
        raise np.linalg.LinAlgError(
            f"propagator rank {generated_rank} is below requested observable rank {requested}; "
            "set truncate=True to restrict to the generated range"
        )
    k = generated_rank if truncate else requested
    if k < 1:
        raise np.linalg.LinAlgError("propagator has empty generated observable range")
    # Preserve observable-POD (decreasing energy) order whenever it is fully
    # generated.  Only rank-deficient truncation rotates to an explicit SVD
    # basis for range(C).
    basis = u[:, :k] if generated_rank < requested else np.eye(requested)
    c = basis.T @ propagator
    generated = ac @ c.T

    # LR: minimum mean-square residual among exact right inverses, with the
    # state covariance as the induced metric (Schlegel's linear regression
    # inverse).  Using covariance with noisy measured b here would destroy
    # C L = I and is therefore deliberately avoided.
    lr_gram = c @ sigma_aa @ c.T
    lr_s = np.linalg.svd(lr_gram, compute_uv=False)
    lr_rank, _ = _numerical_rank(lr_s, lr_gram.shape, relative_rcond)
    if lr_rank < k:
        raise np.linalg.LinAlgError("LR Gram matrix is rank deficient on generated range")
    lr_inverse = sigma_aa @ c.T @ _pinv(lr_gram, relative_rcond)

    g = _as_metric(state_metric, a.shape[1])
    _, _ = _metric_root(g, positive_definite=True)
    g_inverse = scipy.linalg.solve(g, np.eye(g.shape[0]), assume_a="pos", check_finite=True)
    gram = c @ g_inverse @ c.T
    gram_s = np.linalg.svd(gram, compute_uv=False)
    gram_rank, _ = _numerical_rank(gram_s, gram.shape, relative_rcond)
    if gram_rank < k:
        raise np.linalg.LinAlgError("LE Gram matrix is rank deficient on generated range")
    le_inverse = g_inverse @ c.T @ _pinv(gram, relative_rcond)

    def result(inverse: np.ndarray) -> OIDInverse:
        projection = inverse @ c
        residual = ac - generated @ inverse.T
        residual_covariance = residual.T @ residual / n
        residual_observable = residual.T @ generated / n
        eye = np.eye(k)
        return OIDInverse(
            propagator=c, inverse=inverse, projection=projection,
            observable_basis=basis, singular_values=singular_values[:k], rank=k,
            rank_tolerance=rank_tolerance, identity_defect=float(np.linalg.norm(c @ inverse - eye, ord=2)),
            projection_defect=float(np.linalg.norm(projection @ projection - projection, ord=2)),
            propagation_defect=float(np.linalg.norm(c @ projection - c, ord=2)),
            residual_covariance=residual_covariance,
            residual_observable_covariance=residual_observable,
        )

    lr = result(lr_inverse)
    le = result(le_inverse)
    physical = np.eye(a.shape[1]) if state_modes is None else np.asarray(state_modes, dtype=float)
    if physical.ndim != 2 or physical.shape[1] != a.shape[1]:
        raise ValueError("state_modes must map state coefficients to physical space")
    # Columns follow observable-POD energy order before any rank truncation.
    return SchlegelOID(
        observable_pod, propagator, lr, le,
        physical @ lr.inverse, physical @ le.inverse, a.mean(axis=0),
        relative_rcond, requested_observable_rank, requested,
        state_covariance_rank, state_covariance_threshold,
        float(covariance_singular_values[state_covariance_rank - 1]),
    )


def fit_pls_type(pod_coefficients: np.ndarray, observables: np.ndarray) -> dict[str, np.ndarray]:
    a, y = np.asarray(pod_coefficients, dtype=float), np.asarray(observables, dtype=float)
    if a.ndim != 2 or y.ndim != 2 or len(a) != len(y) or len(a) < 2:
        raise ValueError("PLS inputs must be aligned two-dimensional samples")
    a_scaler, y_scaler = Standardizer.fit(a), Standardizer.fit(y)
    a_std, y_std = a_scaler.transform(a), y_scaler.transform(y)
    cross_covariance = a_std.T @ y_std / len(a_std)
    u, singular_values, vt = np.linalg.svd(cross_covariance, full_matrices=False)
    return {"U": u, "S": singular_values, "Vt": vt, "cross_covariance": cross_covariance,
            "a_mean": a_scaler.mean, "a_scale": a_scaler.scale,
            "y_mean": y_scaler.mean, "y_scale": y_scaler.scale,
            "coordinates": a_std @ u}


def physical_pls_modes(pod_modes: np.ndarray, coefficient_scale: np.ndarray, directions: np.ndarray) -> np.ndarray:
    phi = np.asarray(pod_modes)
    scale = np.asarray(coefficient_scale).reshape(-1)
    u = np.asarray(directions)
    if phi.shape[1] != len(scale) or u.shape[0] != len(scale):
        raise ValueError("Phi, D, and U dimensions are inconsistent")
    return phi @ (scale[:, None] * u)


def fit_regularized_cca(a: np.ndarray, y: np.ndarray, epsilon: float) -> dict[str, np.ndarray]:
    x, z = np.asarray(a, dtype=float), np.asarray(y, dtype=float)
    if epsilon <= 0 or x.ndim != 2 or z.ndim != 2 or len(x) != len(z):
        raise ValueError("invalid CCA inputs")
    xs, zs = Standardizer.fit(x).transform(x), Standardizer.fit(z).transform(z)
    n = len(xs)
    cxx = xs.T @ xs / n + epsilon * np.eye(xs.shape[1])
    czz = zs.T @ zs / n + epsilon * np.eye(zs.shape[1])
    cxz = xs.T @ zs / n
    def symmetric_inverse_sqrt(matrix: np.ndarray) -> np.ndarray:
        symmetric = 0.5 * (matrix + matrix.T)
        eigenvalues, eigenvectors = scipy.linalg.eigh(symmetric, check_finite=True)
        tolerance = np.finfo(float).eps * max(matrix.shape) * max(float(eigenvalues[-1]), 1.0)
        if eigenvalues[0] <= tolerance:
            raise np.linalg.LinAlgError("regularized covariance is not positive definite")
        return (eigenvectors * eigenvalues**-0.5) @ eigenvectors.T

    wx = symmetric_inverse_sqrt(cxx)
    wz = symmetric_inverse_sqrt(czz)
    u, s, vt = np.linalg.svd(wx @ cxz @ wz, full_matrices=False)
    return {"W_x": wx @ u, "W_y": wz @ vt.T, "S": s, "epsilon": np.asarray(epsilon)}


def _weighted_inner(a: np.ndarray, b: np.ndarray, weights: np.ndarray | None) -> float:
    if weights is None:
        return float(np.vdot(a, b).real)
    w = np.asarray(weights).reshape(-1)
    return float(np.vdot(a.reshape(-1) * w, b.reshape(-1)).real)


def modal_assurance(a: np.ndarray, b: np.ndarray, weights: np.ndarray | None = None) -> float:
    numerator = _weighted_inner(a, b, weights) ** 2
    denominator = _weighted_inner(a, a, weights) * _weighted_inner(b, b, weights)
    return float(numerator / denominator) if denominator > 0 else 0.0


def principal_angles(a: np.ndarray, b: np.ndarray, weights: np.ndarray | None = None) -> np.ndarray:
    x, y = np.asarray(a, dtype=float), np.asarray(b, dtype=float)
    if x.ndim != 2 or y.ndim != 2 or x.shape[0] != y.shape[0]:
        raise ValueError("subspace bases must share physical dimension")
    if weights is not None:
        root = np.sqrt(np.asarray(weights).reshape(-1, 1))
        x, y = root * x, root * y
    qx = np.linalg.qr(x)[0]
    qy = np.linalg.qr(y)[0]
    return scipy.linalg.subspace_angles(qx, qy)


def projector_similarity(a: np.ndarray, b: np.ndarray, weights: np.ndarray | None = None) -> float:
    angles = principal_angles(a, b, weights)
    return float(np.mean(np.cos(angles) ** 2))
