"""Synthetic finite-dimensional contracts for Schlegel LR/LE-OID."""
from __future__ import annotations

import numpy as np
import pytest

from OID_analysis.v2.analysis import fit_observable_pod, fit_pod, fit_schlegel_oid


def _problem(seed=12, n=2000):
    rng = np.random.default_rng(seed)
    mixing = rng.normal(size=(5, 5))
    covariance = mixing @ mixing.T + np.eye(5)
    a = rng.multivariate_normal(np.zeros(5), covariance, size=n)
    c = np.array([[1.0, -0.4, 0.2, 0.0, 0.5],
                  [0.1, 0.3, -0.7, 0.8, 0.0]])
    b = a @ c.T
    return a, b


def test_weighted_state_pod_respects_diagonal_fluid_mask():
    rng = np.random.default_rng(1)
    q = rng.normal(size=(300, 6))
    q[:, 5] = rng.normal(scale=1e6, size=300)  # solid-cell contamination
    weights = np.array([2.0, 0.5, 3.0, 1.0, 4.0, 0.0])
    pod = fit_pod(q, 5, weights)
    assert np.allclose(pod.modes.T @ np.diag(weights) @ pod.modes, np.eye(5))
    assert np.allclose(pod.coefficients, (q - pod.mean) @ np.diag(weights) @ pod.modes)
    assert np.allclose(pod.inverse_transform(pod.coefficients)[:, :5], q[:, :5], atol=1e-10)


def test_observable_metric_is_equivalent_to_euclidean_transformation():
    rng = np.random.default_rng(2)
    y = rng.normal(size=(500, 3))
    g = np.diag([4.0, 0.25, 2.0])
    pod = fit_observable_pod(y, 3, g)
    root = np.diag(np.sqrt(np.diag(g)))
    transformed = fit_observable_pod(y @ root, 3)
    assert np.allclose(pod.singular_values, transformed.singular_values)
    assert np.allclose(np.abs(pod.modes.T @ g @ pod.modes), np.eye(3), atol=1e-10)
    assert np.allclose(np.abs(pod.coefficients.T @ transformed.coefficients),
                       np.diag(pod.singular_values**2), atol=1e-8)


def test_lr_le_right_inverse_projection_and_propagation_identities():
    a, b = _problem()
    oid = fit_schlegel_oid(a, b, observable_rank=2)
    for inverse in (oid.lr, oid.le):
        c, l, p = inverse.propagator, inverse.inverse, inverse.projection
        assert np.allclose(c @ l, np.eye(2), atol=1e-10)
        assert np.allclose(p @ p, p, atol=1e-10)
        assert np.allclose(c @ p, c, atol=1e-10)
        assert inverse.identity_defect < 1e-10
        assert inverse.projection_defect < 1e-10
        assert inverse.propagation_defect < 1e-10


def test_lr_residual_is_uncorrelated_and_variance_optimal():
    a, b = _problem(seed=4)
    oid = fit_schlegel_oid(a, b, observable_rank=2)
    lr = oid.lr
    assert np.linalg.norm(lr.residual_observable_covariance) < 1e-10
    rng = np.random.default_rng(5)
    null = scipy_null(lr.propagator)
    perturbation = null @ rng.normal(size=(null.shape[1], lr.rank))
    alternative = lr.inverse + perturbation
    centered = a - a.mean(axis=0)
    generated = centered @ lr.propagator.T
    residual_alt = centered - generated @ alternative.T
    assert np.sum(residual_alt**2) >= np.sum(lr.residual_covariance) * len(a) - 1e-7


def scipy_null(c):
    _, s, vt = np.linalg.svd(c, full_matrices=True)
    rank = np.sum(s > np.finfo(float).eps * max(c.shape) * s[0])
    return vt[rank:].T


def test_le_is_g_orthogonal_and_minimum_energy_under_null_perturbations():
    a, b = _problem(seed=7)
    g = np.diag([1.0, 2.0, 4.0, 0.5, 3.0])
    le = fit_schlegel_oid(a, b, observable_rank=2, state_metric=g).le
    null = scipy_null(le.propagator)
    assert np.allclose(null.T @ g @ le.inverse, 0.0, atol=1e-10)
    rng = np.random.default_rng(8)
    perturbation = null @ rng.normal(size=(null.shape[1], le.rank))
    alternative = le.inverse + perturbation
    energy = lambda x: np.trace(x.T @ g @ x)
    assert energy(alternative) >= energy(le.inverse) - 1e-10


def test_oid_modes_follow_observable_pod_energy_order():
    rng = np.random.default_rng(9)
    a = rng.normal(size=(1000, 4))
    y = np.column_stack((10.0 * a[:, 2], a[:, 0]))
    physical = np.linalg.qr(rng.normal(size=(12, 4)))[0]
    oid = fit_schlegel_oid(a, y, observable_rank=2, state_modes=physical)
    assert oid.observable_pod.singular_values[0] > oid.observable_pod.singular_values[1]
    assert np.allclose(oid.lr_modes, physical @ oid.lr.inverse)
    assert np.allclose(oid.le_modes, physical @ oid.le.inverse)


def test_rank_deficiency_fails_closed_or_truncates_generated_range():
    rng = np.random.default_rng(10)
    a = rng.normal(size=(500, 4))
    noise = rng.normal(size=len(a))
    y = np.column_stack((a[:, 0], 2.0 * a[:, 0] + noise, a[:, 1]))
    with pytest.raises(np.linalg.LinAlgError, match="set truncate=True"):
        fit_schlegel_oid(a, y, observable_rank=3, rcond=0.05)
    oid = fit_schlegel_oid(a, y, observable_rank=3, rcond=0.05, truncate=True)
    assert oid.lr.rank == 2
    assert oid.le.rank == 2
    assert np.allclose(oid.lr.propagator @ oid.lr.inverse, np.eye(2), atol=1e-10)


def test_propagator_is_regression_not_cross_covariance():
    rng = np.random.default_rng(11)
    a = rng.normal(size=(2000, 3)) @ np.diag([1.0, 3.0, 7.0])
    true_c = np.array([[2.0, -1.0, 0.5]])
    y = a @ true_c.T
    oid = fit_schlegel_oid(a, y, observable_rank=1)
    # Observable POD can flip sign, but regression must recover C up to that sign.
    assert np.allclose(np.abs(oid.propagator), np.abs(true_c), rtol=1e-10, atol=1e-10)
    cross = (oid.observable_pod.coefficients.T @ (a - a.mean(axis=0))) / len(a)
    assert not np.allclose(np.abs(cross), np.abs(oid.propagator))


def test_snapshot_pod_matches_direct_weighted_svd_and_reconstructs():
    rng = np.random.default_rng(21)
    data = rng.normal(size=(18, 73))
    weights = rng.uniform(0.1, 3.0, size=data.shape[1])
    pod = fit_pod(data, rank=17, weights=weights)
    centered = data - data.mean(axis=0)
    _, singular_values, vt = np.linalg.svd(centered * np.sqrt(weights), full_matrices=False)
    direct_modes = vt[:17].T / np.sqrt(weights)[:, None]
    gram = (centered * np.sqrt(weights)) @ (centered * np.sqrt(weights)).T
    eigenvalues, eigenvectors = np.linalg.eigh(gram)
    order = np.argsort(eigenvalues)[::-1][:17]
    snapshot_modes = centered.T @ eigenvectors[:, order] / np.sqrt(eigenvalues[order])

    assert np.allclose(pod.singular_values, singular_values[:17], rtol=1e-11, atol=1e-11)
    assert np.allclose(np.abs(pod.modes.T @ (weights[:, None] * direct_modes)),
                       np.eye(17), atol=1e-9)
    assert np.allclose(np.abs(pod.modes.T @ (weights[:, None] * snapshot_modes)),
                       np.eye(17), atol=1e-9)
    assert np.allclose(pod.modes.T @ (weights[:, None] * pod.modes),
                       np.eye(17), atol=1e-10)
    assert np.allclose(pod.transform(data), pod.coefficients, atol=1e-10)
    assert np.allclose(pod.inverse_transform(pod.coefficients), data, atol=1e-10)
    assert np.isclose(pod.total_energy, np.sum(singular_values**2), rtol=1e-12)


def test_snapshot_pod_zero_weights_remove_contamination_and_zero_modes():
    rng = np.random.default_rng(22)
    clean = rng.normal(size=(14, 20))
    weights = np.r_[np.ones(15), np.zeros(5)]
    contaminated = clean.copy()
    contaminated[:, -5:] = rng.normal(scale=1e100, size=(14, 5))
    first = fit_pod(clean, rank=8, weights=weights)
    second = fit_pod(contaminated, rank=8, weights=weights)

    assert np.array_equal(first.modes[-5:], np.zeros((5, 8)))
    assert np.array_equal(second.modes[-5:], np.zeros((5, 8)))
    assert np.allclose(first.singular_values, second.singular_values)
    assert np.allclose(first.coefficients, second.coefficients)
    assert np.allclose(second.inverse_transform(second.coefficients)[:, :15],
                       contaminated.mean(axis=0)[:15] +
                       second.coefficients @ second.modes[:15].T)


def test_field_pod_never_allocates_diagonal_metric(monkeypatch):
    rng = np.random.default_rng(23)
    data = rng.normal(size=(9, 100_000))
    weights = np.ones(data.shape[1])

    def forbidden(*args, **kwargs):
        raise AssertionError("field POD must not call np.diag")

    monkeypatch.setattr(np, "diag", forbidden)
    pod = fit_pod(data, rank=4, weights=weights)
    assert pod.metric.ndim == 1
    assert pod.modes.shape == (100_000, 4)


def test_rank_max_is_sliced_to_weighted_numerical_rank():
    rng = np.random.default_rng(24)
    data = rng.normal(size=(7, 30))
    pod = fit_pod(data, rank=100)
    assert pod.modes.shape[1] == 6  # centering removes one snapshot direction
    assert np.allclose(pod.inverse_transform(pod.transform(data)), data, atol=1e-10)


def test_generated_range_prediction_and_rcond_zero():
    rng = np.random.default_rng(25)
    a = rng.normal(size=(300, 4))
    y = np.column_stack((a[:, 0], 2.0 * a[:, 0], a[:, 1]))
    oid = fit_schlegel_oid(a, y, observable_rank=3, truncate=True, rcond=0)

    assert oid.lr.rank == 2
    assert oid.lr.rank_tolerance == 0.0
    assert oid.observable_coordinates(y).shape == (len(y), 2)
    assert oid.predict_observables(a).shape == y.shape
    assert np.allclose(oid.predict_state_coefficients(y) - oid.state_mean,
                       oid.observable_coordinates(y) @ oid.lr.inverse.T)


def test_one_rcond_controls_covariance_propagator_and_inverse_ranks():
    rng = np.random.default_rng(26)
    base = rng.normal(size=(1000, 3))
    a = base * np.array([1.0, 1e-5, 1e-8])
    y = np.column_stack((a[:, 0] + a[:, 1], a[:, 0] - a[:, 1]))
    oid = fit_schlegel_oid(
        a, y, observable_rank=2, truncate=True, rcond=1e-12
    )

    assert oid.rcond == 1e-12
    assert oid.state_covariance_rank == 2
    covariance_s = np.linalg.svd(np.cov(a, rowvar=False, bias=True), compute_uv=False)
    assert oid.state_covariance_threshold == pytest.approx(1e-12 * covariance_s[0])
    assert oid.state_covariance_smallest_retained_singular == pytest.approx(covariance_s[1])
    assert oid.lr.rank == oid.le.rank == 2
    assert oid.lr.rank_tolerance == pytest.approx(
        1e-12 * np.linalg.svd(oid.propagator, compute_uv=False)[0]
    )
    assert np.allclose(oid.lr.propagator @ oid.lr.inverse, np.eye(2), atol=1e-8)
    assert np.allclose(oid.le.propagator @ oid.le.inverse, np.eye(2), atol=1e-8)


def test_diagonal_state_metric_matches_dense_diagonal_le_inverse():
    a, b = _problem(seed=31, n=1200)
    weights = np.array([1.0, 2.0, 4.0, 0.5, 3.0])
    vector = fit_schlegel_oid(a, b, observable_rank=2, state_metric=weights)
    dense = fit_schlegel_oid(a, b, observable_rank=2, state_metric=np.diag(weights))
    assert np.allclose(vector.le.inverse, dense.le.inverse, rtol=1e-11, atol=1e-11)
    assert np.allclose(vector.le.projection, dense.le.projection, rtol=1e-11, atol=1e-11)
