from __future__ import annotations

import importlib
import os
import sys
from pathlib import Path

import numpy as np
import pytest

# Optional: allow running tests from source tree without installing the package.
# Default is OFF so that CI (and users) test the installed distribution (wheel).
#
# Usage:
#   Linux/macOS: ZEN_FRONTS_TEST_FROM_SRC=1 pytest -q
#   Windows (PowerShell): $env:ZEN_FRONTS_TEST_FROM_SRC="1"; pytest -q
if os.environ.get("ZEN_FRONTS_TEST_FROM_SRC") in {"1", "true", "True", "yes", "YES"}:
    ROOT = Path(__file__).resolve().parents[1]
    SRC = ROOT / "src"
    if str(SRC) not in sys.path:
        sys.path.insert(0, str(SRC))

from zen_fronts.pareto_core.pareto_rating import (  # noqa: E402
    ParetoRating,
    crowding_distance,
    non_dominated_sort,
    rate_population,
    transform_to_ranks,
)


def test_non_dominated_sort_basic_two_objectives() -> None:
    # Maximize. Expect 3 fronts.
    X = np.array(
        [
            [1.0, 1.0],  # A
            [0.8, 0.8],  # B
            [1.0, 0.0],  # C
            [0.0, 1.0],  # D
            [0.2, 0.2],  # E
        ],
        dtype=np.float64,
    )

    front_id, fronts = non_dominated_sort(X)

    assert front_id.tolist() == [0, 1, 1, 1, 2]
    assert [f.tolist() for f in fronts] == [[0], [1, 2, 3], [4]]


def test_crowding_distance_interior_values() -> None:
    # 4-point non-dominated front in 2D.
    # Points lie on a decreasing curve: extremes should get +inf.
    X = np.array(
        [
            [0.0, 1.0],
            [0.3, 0.7],
            [0.7, 0.3],
            [1.0, 0.0],
        ],
        dtype=np.float64,
    )
    front = np.array([0, 1, 2, 3], dtype=np.int64)

    out = np.zeros(X.shape[0], dtype=np.float64)
    crowding_distance(X, front, out)

    assert np.isinf(out[0])
    assert np.isinf(out[3])
    # Both interior points have symmetric crowding.
    assert out[1] == pytest.approx(1.4)
    assert out[2] == pytest.approx(1.4)

    # Sorting within a single front must be deterministic.
    rating = rate_population(X, mode="values", within_mode="crowding")
    assert rating.order.tolist() == [0, 3, 1, 2]


def test_crowding_distance_degenerate_objective_does_not_create_fake_extremes() -> None:
    # Second objective is constant: no +inf should be created because of it.
    X = np.array(
        [
            [0.0, 0.0],
            [0.5, 0.0],
            [1.0, 0.0],
        ],
        dtype=np.float64,
    )
    front = np.array([0, 1, 2], dtype=np.int64)
    out = np.zeros(X.shape[0], dtype=np.float64)

    crowding_distance(X, front, out)

    assert np.isinf(out[0])
    assert np.isinf(out[2])
    assert out[1] == pytest.approx(1.0)


def test_transform_to_ranks_is_stable_by_index_for_ties() -> None:
    # All values tied: ranks should be assigned by index order.
    X = np.array(
        [
            [1.0, 1.0],
            [1.0, 1.0],
            [1.0, 1.0],
            [1.0, 1.0],
        ],
        dtype=np.float64,
    )

    R = transform_to_ranks(X)

    # Positions 0..N-1 normalized to [0,1].
    expected = np.array([0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0], dtype=np.float64)
    assert np.allclose(R[:, 0], expected)
    assert np.allclose(R[:, 1], expected)


def test_rate_population_outputs_are_deterministic_and_tie_breaks_by_index() -> None:
    X = np.array(
        [
            [1.0, 1.0],  # A
            [0.8, 0.8],  # B
            [1.0, 0.0],  # C
            [0.0, 1.0],  # D
            [0.2, 0.2],  # E
        ],
        dtype=np.float64,
    )

    r1 = rate_population(X, mode="values", within_mode="crowding")
    r2 = rate_population(X, mode="values", within_mode="crowding")

    assert isinstance(r1, ParetoRating)
    assert np.array_equal(r1.front_id, r2.front_id)
    assert np.array_equal(r1.order, r2.order)

    # Expected deterministic order:
    # front0: [0]
    # front1: extremes [2,3] (both +inf), then [1]
    # front2: [4]
    assert r1.order.tolist() == [0, 2, 3, 1, 4]


# --------------------------- Additional tests ---------------------------


@pytest.mark.parametrize("use_cython", [False, True])
def test_rate_population_empty_and_singleton(use_cython: bool) -> None:
    X0 = np.zeros((0, 3), dtype=np.float64)
    r0 = rate_population(
        X0, mode="values", within_mode="crowding", use_cython=use_cython
    )
    assert r0.front_id.shape == (0,)
    assert r0.within_score.shape == (0,)
    assert r0.order.shape == (0,)
    assert r0.fronts == []

    X1 = np.array([[0.1, 0.2]], dtype=np.float64)
    r1 = rate_population(
        X1, mode="values", within_mode="crowding", use_cython=use_cython
    )
    assert r1.front_id.tolist() == [0]
    assert np.isinf(r1.within_score[0])
    assert r1.order.tolist() == [0]
    assert [f.tolist() for f in r1.fronts] == [[0]]


@pytest.mark.parametrize("use_cython", [False, True])
def test_one_objective_fronts_are_layers_by_value(use_cython: bool) -> None:
    # 1D maximize:
    # value 1.0 dominates 0.5 dominates 0.0 (ties are non-dominating among themselves).
    X = np.array([[0.0], [0.0], [1.0], [0.5]], dtype=np.float64)
    r = rate_population(X, mode="values", within_mode="crowding", use_cython=use_cython)

    assert r.front_id.tolist() == [2, 2, 0, 1]
    assert r.order.tolist() == [2, 3, 0, 1]


@pytest.mark.parametrize("use_cython", [False, True])
def test_duplicates_do_not_dominate_each_other_and_tie_break_by_index(
    use_cython: bool,
) -> None:
    # Two identical best points and one dominated point.
    X = np.array(
        [
            [1.0, 1.0],  # 0
            [1.0, 1.0],  # 1 (duplicate)
            [0.0, 0.0],  # 2 dominated by both
        ],
        dtype=np.float64,
    )
    r = rate_population(X, mode="values", within_mode="crowding", use_cython=use_cython)

    # Front0 has [0,1] (sorted by index), front1 has [2].
    assert [f.tolist() for f in r.fronts] == [[0, 1], [2]]
    assert r.front_id.tolist() == [0, 0, 1]

    # For front0: all objectives degenerate within that front => crowding all zeros.
    assert r.within_score[0] == pytest.approx(0.0)
    assert r.within_score[1] == pytest.approx(0.0)

    # Global order tie-breaks by index.
    assert r.order.tolist() == [0, 1, 2]


def test_crowding_two_points_non_degenerate_gives_infinite_for_both() -> None:
    X = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.float64)
    front = np.array([0, 1], dtype=np.int64)
    out = np.zeros(2, dtype=np.float64)

    crowding_distance(X, front, out)
    assert np.isinf(out[0])
    assert np.isinf(out[1])


def test_crowding_all_objectives_degenerate_gives_zero_and_stable_order() -> None:
    # All points identical => single front, but no objective has spread => crowding all zeros.
    X = np.array([[0.5, 0.5], [0.5, 0.5], [0.5, 0.5]], dtype=np.float64)
    r = rate_population(X, mode="values", within_mode="crowding", use_cython=False)

    assert [f.tolist() for f in r.fronts] == [[0, 1, 2]]
    assert np.allclose(r.within_score, 0.0)
    assert r.order.tolist() == [0, 1, 2]


@pytest.mark.parametrize("use_cython", [False, True])
def test_rank_mode_equivalent_to_explicit_transform(use_cython: bool) -> None:
    X = np.array(
        [
            [10.0, 0.0, 5.0],
            [9.0, 1.0, 5.0],
            [0.0, 10.0, 4.0],
            [9.0, 1.0, 6.0],
            [9.0, 1.0, 6.0],
        ],
        dtype=np.float64,
    )
    R = transform_to_ranks(X)

    r1 = rate_population(X, mode="ranks", within_mode="crowding", use_cython=use_cython)
    r2 = rate_population(
        R, mode="values", within_mode="crowding", use_cython=use_cython
    )

    assert np.array_equal(r1.front_id, r2.front_id)
    assert np.array_equal(r1.order, r2.order)
    assert np.allclose(r1.within_score, r2.within_score, equal_nan=True)


@pytest.mark.parametrize("use_cython", [False, True])
def test_monotone_dominance_chain_fronts(use_cython: bool) -> None:
    # Strict chain: i dominates i+1 for all i
    X = np.array(
        [
            [1.0, 1.0, 1.0],
            [0.9, 0.9, 0.9],
            [0.8, 0.8, 0.8],
            [0.7, 0.7, 0.7],
        ],
        dtype=np.float64,
    )
    r = rate_population(X, mode="values", within_mode="crowding", use_cython=use_cython)
    assert r.front_id.tolist() == [0, 1, 2, 3]
    assert r.order.tolist() == [0, 1, 2, 3]


@pytest.mark.parametrize("use_cython", [False, True])
def test_property_dominance_implies_strictly_better_front(use_cython: bool) -> None:
    rng = np.random.default_rng(0)
    X = rng.random((50, 4), dtype=np.float64)

    r = rate_population(X, mode="values", within_mode="crowding", use_cython=use_cython)
    fid = r.front_id

    ge = X[:, None, :] >= X[None, :, :]
    gt = X[:, None, :] > X[None, :, :]
    dom = np.all(ge, axis=2) & np.any(gt, axis=2)
    np.fill_diagonal(dom, False)

    # If i dominates j, then i must be in a strictly better front.
    viol = dom & ~(fid[:, None] < fid[None, :])
    assert not viol.any()


def _cython_ext_available() -> bool:
    try:
        importlib.import_module("zen_fronts.pareto_core._pareto_rating_cy")
        return True
    except Exception:
        return False


@pytest.mark.skipif(
    not _cython_ext_available(),
    reason="Cython extension is not built (zen_fronts.pareto_core._pareto_rating_cy)",
)
def test_cython_matches_python_fronts_and_order_on_fixed_and_random_cases() -> None:
    rng = np.random.default_rng(123)

    fixed_cases = [
        np.array([[0.1, 0.2]], dtype=np.float64),
        np.array(
            [[1.0, 1.0], [1.0, 1.0], [0.0, 0.0]], dtype=np.float64
        ),  # duplicates + dominated
        np.array(
            [[0.0, 0.0], [0.5, 0.0], [1.0, 0.0]], dtype=np.float64
        ),  # degenerate objective
        np.array(
            [[0.0, 1.0], [0.3, 0.7], [0.7, 0.3], [1.0, 0.0]], dtype=np.float64
        ),  # classic
        np.array([[1.0], [0.0], [0.0], [0.5]], dtype=np.float64),  # 1D with ties
        np.array(  # includes ties in some columns
            [
                [1.0, 2.0, 3.0],
                [1.0, 2.0, 2.5],
                [0.9, 2.0, 3.0],
                [1.0, 1.5, 3.0],
            ],
            dtype=np.float64,
        ),
    ]

    random_cases = [
        rng.random((50, 4), dtype=np.float64),
        rng.random((100, 3), dtype=np.float64),
    ]

    # Add a random case with many ties (to stress rank-mode determinism).
    X_ties = rng.integers(0, 5, size=(80, 4)).astype(np.float64) / 4.0
    random_cases.append(X_ties)

    for mode in ("values", "ranks"):
        for X in fixed_cases + random_cases:
            r_py = rate_population(
                X, mode=mode, within_mode="crowding", use_cython=False
            )
            r_cy = rate_population(
                X, mode=mode, within_mode="crowding", use_cython=True
            )

            # Hard requirements: fronts and order must match exactly.
            assert np.array_equal(r_py.front_id, r_cy.front_id)
            assert np.array_equal(r_py.order, r_cy.order)

            # Front partition representation should match as well.
            assert len(r_py.fronts) == len(r_cy.fronts)
            for f_py, f_cy in zip(r_py.fronts, r_cy.fronts):
                assert np.array_equal(f_py, f_cy)

            # within_score: exact agreement on inf positions; tight numeric match on finite scores.
            inf_py = np.isinf(r_py.within_score)
            inf_cy = np.isinf(r_cy.within_score)
            assert np.array_equal(inf_py, inf_cy)

            finite = ~inf_py
            assert np.allclose(
                r_py.within_score[finite],
                r_cy.within_score[finite],
                rtol=0.0,
                atol=1e-12,
            )
