from __future__ import annotations

import numpy as np

from zen_fronts.crits import CriteriaStore, CriteriaStoreConfig


def test_active_ids_only_when_all_criteria_ready() -> None:
    store = CriteriaStore({"a": "max", "b": "min"}, config=CriteriaStoreConfig(project_trend=False))

    store.register_point(0, active=True)
    store.register_point(1, active=True)

    # No updates => none active
    assert store.active_point_ids() == []

    store.update_crits(0, {"a": {"mu": 1.0, "sigma2": 0.0}}, t=1.0)
    assert store.active_point_ids() == []

    store.update_crits(0, {"b": {"mu": 2.0, "sigma2": 0.0}}, t=1.0)
    assert store.active_point_ids() == [0]

    store.update_crits(1, {"a": {"mu": 10.0, "sigma2": 0.0}, "b": {"mu": -5.0, "sigma2": 0.0}}, t=1.0)
    assert store.active_point_ids() == [0, 1]


def test_build_matrices_projects_trend_and_flips_min_objectives() -> None:
    store = CriteriaStore({"a": "max", "b": "min"}, config=CriteriaStoreConfig(project_trend=True))
    store.register_point(0, active=True)

    store.update_crits(
        0,
        {
            "a": {"mu": 1.0, "trend": 2.0, "sigma2": 0.5},
            "b": {"mu": 3.0, "trend": -1.0, "sigma2": 1.5},
        },
        t=10.0,
    )

    mu, sigma2 = store.build_matrices([0], now=12.0)
    assert mu.shape == (1, 2)
    assert sigma2.shape == (1, 2)

    # a: 1 + 2*(12-10) = 5
    assert float(mu[0, 0]) == 5.0
    assert float(sigma2[0, 0]) == 0.5

    # b is MIN => sign flipped after projection.
    # b proj: 3 + (-1)*(2) = 1  -> flip => -1
    assert float(mu[0, 1]) == -1.0
    assert float(sigma2[0, 1]) == 1.5


def test_raw_and_stats_paths_can_match_matrix_values() -> None:
    """We do not assume upstream `owrf1d` behavior.

    Instead, we assert a weaker but valuable property: if external stats are
    set to the state's current values produced by the raw path, the matrices
    match.
    """

    store = CriteriaStore({"c": "max"})
    store.register_point(0, active=True)
    store.register_point(1, active=True)

    # Raw updates for point 0
    store.update_crits(0, {"c": 1.0}, t=1.0)
    store.update_crits(0, {"c": 2.0}, t=2.0)
    st = store.state(0, "c")
    assert st is not None and st.ready

    # Apply equivalent external stats to point 1
    store.update_crits(
        1,
        {"c": {"mu": float(st.mu), "trend": float(st.trend), "sigma2": float(st.sigma2)}},
        t=float(st.t_last),
    )

    mu, sigma2 = store.build_matrices([0, 1], now=float(st.t_last))
    assert np.allclose(mu[0], mu[1])
    assert np.allclose(sigma2[0], sigma2[1])
