import numpy as np

from measurekit.backends.numpy_backend import NumpyBackend
from measurekit.domain.measurement.quantity import Quantity
from measurekit.domain.measurement.units import (
    get_default_system,
)
from measurekit.domain.measurement.vectorized_uncertainty import (
    ensure_store,
)


def test_correlated_arrays():
    system = get_default_system()
    meter = system.get_unit("meter")
    store = ensure_store(NumpyBackend())

    val_base = np.array([1.0, 2.0])
    unc_base = np.array([0.1, 0.2])
    q_base = Quantity.from_input(val_base, meter, system, uncertainty=unc_base)

    # q1 = q_base + 10
    q1 = q_base + Quantity.from_input(
        np.array([10.0, 10.0]), meter, system, uncertainty=0.0
    )
    # q2 = q_base * 2
    q2 = q_base * 2.0

    # q3 = q1 + q2 = 3 * q_base + 10
    q3 = q1 + q2

    print(f"q1 uncertainty: {q1.uncertainty}")
    print(f"q2 uncertainty: {q2.uncertainty}")
    print(f"q3 uncertainty: {q3.uncertainty}")

    expected_unc = 3.0 * unc_base
    # The following assertion requires implicit cross-correlation which is pending Phase 4.
    # We test it below in a try/except block.
    # np.testing.assert_allclose(q3.uncertainty, expected_unc)

    # Verify cross-correlation: Cov(q1, q2) = 2*Var(q_base)
    s1 = q1.uncertainty_obj.vector_slice
    s2 = q2.uncertainty_obj.vector_slice

    cov_12 = store.get_covariance_block(s1, s2).toarray()
    expected_cov_12 = 2.0 * np.diag(unc_base**2)
    np.testing.assert_allclose(
        cov_12, expected_cov_12, err_msg="Cross-covariance mismatch!"
    )
    print("Success: Cross-cov matches theoretical 2 * Var(q_base).")

    # Verify q3 = q1 + q2 via full covariance:
    # Var(q3) = Var(q1) + Var(q2) + 2*Cov(q1, q2)
    np.testing.assert_allclose(
        q3.uncertainty,
        expected_unc,
        err_msg="Implicit cross-correlation failed!",
    )
    print("Success: q3 uncertainty matches theoretical expectation.")

    # Verify Root Propagation (Sanity Check for Rust Backend)
    # q4 = 3 * q_base + 10 meters
    offset = Quantity.from_input(
        np.array([10.0, 10.0]), meter, system, uncertainty=0.0
    )
    q4 = 3.0 * q_base + offset
    np.testing.assert_allclose(q4.uncertainty, expected_unc)
    print(
        "Success: Root propagation matches theoretical 3 * q_base uncertainty."
    )


if __name__ == "__main__":
    test_correlated_arrays()
