"""Test asdex sparsity detection against CUTEst ground truth via sif2jax.

These tests compare asdex's detected sparsity patterns against
ground truth from CUTEst benchmark problems.
The goal is to catalog conservativeness — where asdex over-estimates sparsity.

Fixtures are generated by ``tests/setup/generate_cutest_fixtures.py``.

sif2jax problem objects have the following attributes:

- ``name``: str — problem name (e.g. "AKIVA").
- ``y0``: jax.Array — starting point, shape ``(n,)``.
- ``args``: None | tuple of arrays — extra data passed to the objective.
- ``objective(y, args) -> scalar`` — the objective function.
- ``constraint(y) -> (eq_or_none, ineq_or_none)`` — equality and inequality
  constraint vectors (constrained problems only).
"""

from __future__ import annotations

import warnings
from pathlib import Path

import numpy as np
import pytest
import sif2jax

import asdex

FIXTURE_DIR = Path(__file__).parent / "cutest_fixtures"

# Expected (detected_nnz, target_nnz) per problem.
# When a handler improvement reduces detected nnz,
# update the corresponding value here.

EXPECTED_HESSIAN_NNZ: dict[str, tuple[int, int]] = {
    "AKIVA": (4, 4),
    "ALLINITU": (16, 6),
    "ARGLINA": (40000, 200),
    "ARGLINB": (40000, 40000),
    "ARGLINC": (39204, 39204),
    "ARGTRIGLS": (40000, 40000),
    "BARD": (9, 9),
    "BEALE": (4, 3),
    "BENNETT5LS": (9, 9),
    "BIGGS6": (36, 36),
    "BOX3": (9, 9),
    "BOXBODLS": (4, 4),
    "BROWNBS": (4, 2),
    "BROWNDEN": (16, 16),
    "CHNROSNB": (148, 148),
    "CHNRSNBM": (148, 148),
    "CHWIRUT1LS": (9, 9),
    "CHWIRUT2LS": (9, 9),
    "CLIFF": (4, 4),
    "CLUSTERLS": (4, 2),
    "COATING": (2300, 2300),
    "COOLHANSLS": (81, 39),
    "CUBE": (4, 4),
    "DANIWOODLS": (4, 4),
    "DENSCHNA": (4, 4),
    "DENSCHNB": (4, 4),
    "DENSCHNC": (4, 4),
    "DENSCHND": (9, 9),
    "DENSCHNE": (3, 3),
    "DENSCHNF": (4, 4),
    "DEVGLA1": (16, 16),
    "DEVGLA2": (25, 25),
    "DIXMAANA1": (9, 9),
    "DJTL": (4, 4),
    "DMN15102LS": (4356, 4356),
    "DMN15103LS": (9801, 9801),
    "EGGCRATE": (2, 2),
    "ELATVIDU": (4, 4),
    "ENGVAL2": (9, 9),
    "ERRINROS": (148, 148),
    "EXP2": (4, 4),
    "EXPFIT": (4, 3),
    "GAUSS1LS": (64, 64),
    "GAUSS2LS": (64, 64),
    "GAUSS3LS": (64, 64),
    "GAUSSIAN": (9, 5),
    "GENROSE": (1498, 1498),
    "GROWTHLS": (9, 9),
    "HAHN1LS": (49, 49),
    "HAIRY": (4, 4),
    "HATFLDD": (9, 9),
    "HATFLDE": (9, 9),
    "HATFLDFL": (9, 9),
    "HATFLDFLS": (9, 9),
    "HEART6LS": (36, 20),
    "HEART8LS": (64, 32),
    "HELIX": (9, 7),
    "HILBERTA": (4, 4),
    "HILBERTB": (100, 100),
    "HIMMELBCLS": (4, 4),
    "HIMMELBG": (4, 4),
    "HIMMELBH": (2, 1),
    "HUMPS": (4, 4),
    "JENSMP": (4, 4),
    "JUDGE": (4, 4),
    "KIRBY2LS": (25, 25),
    "KOWOSB": (16, 16),
    "LANCZOS1LS": (36, 36),
    "LANCZOS2LS": (36, 36),
    "LOGHAIRY": (4, 4),
    "LSC1LS": (9, 9),
    "LSC2LS": (9, 9),
    "LUKSAN11LS": (298, 298),
    "LUKSAN12LS": (484, 484),
    "LUKSAN13LS": (484, 482),
    "LUKSAN14LS": (420, 420),
    "LUKSAN15LS": (592, 592),
    "LUKSAN16LS": (592, 592),
    "LUKSAN17LS": (592, 592),
    "LUKSAN21LS": (494, 494),
    "MARATOSB": (4, 4),
    "MEXHAT": (4, 4),
    "MGH09LS": (16, 16),
    "MGH10LS": (9, 9),
    "MGH10SLS": (9, 9),
    "MGH17LS": (25, 25),
    "MGH17SLS": (25, 19),
    "MISRA1ALS": (4, 4),
    "MISRA1BLS": (4, 4),
    "MISRA1CLS": (4, 4),
    "MISRA1DLS": (4, 4),
    "NELSONLS": (9, 9),
    "OSBORNEA": (25, 25),
    "OSBORNEB": (121, 117),
    "PALMER1C": (64, 64),
    "PALMER1D": (49, 49),
    "PALMER2C": (64, 64),
    "PALMER3C": (64, 64),
    "PALMER4C": (64, 64),
    "PALMER5C": (36, 36),
    "PALMER5D": (16, 16),
    "PALMER6C": (64, 64),
    "PALMER7C": (64, 64),
    "PALMER8C": (64, 64),
    "PENALTY3": (40000, 40000),
    "PRICE3": (4, 4),
    "PRICE4": (4, 4),
    "QING": (100, 99),
    "RAT42LS": (9, 9),
    "RAT43LS": (16, 16),
    "ROSENBR": (4, 4),
    "ROSZMAN1LS": (16, 16),
    "S308": (4, 4),
    "SISSER": (4, 4),
    "SNAIL": (4, 4),
    "SPIN2LS": (10404, 10389),
    "TRIGON1": (100, 100),
    "VANDANMSLS": (484, 16),
    "VARDIM": (40000, 40000),
    "VESUVIALS": (64, 46),
    "VESUVIOLS": (64, 46),
    "VESUVIOULS": (64, 42),
    "VIBRBEAM": (64, 64),
    "WAYSEA1": (4, 4),
    "WAYSEA2": (4, 4),
    "ZANGWIL2": (4, 4),
}

EXPECTED_JACOBIAN_EQ_NNZ: dict[str, tuple[int, int]] = {
    "ALSOTAME": (2, 2),
    "BATCH": (54, 48),
    "BT10": (4, 4),
    "BT11": (8, 8),
    "BT12": (8, 7),
    "BT3": (7, 7),
    "BT4": (6, 6),
    "BT5": (6, 6),
    "BT6": (5, 5),
    "BT7": (8, 8),
    "BT8": (6, 4),
    "BT9": (6, 6),
    "BYRDSPHR": (6, 6),
    "CSFI1": (7, 7),
    "CSFI2": (7, 7),
    "DALLASS": (92, 92),
    "DECONVC": (11, 11),
    "DEGENLPA": (82, 82),
    "DEGENLPB": (82, 82),
    "DEGENQP": (10, 10),
    "DEGENQPC": (50, 50),
    "DUAL1": (85, 85),
    "DUAL2": (96, 96),
    "DUAL3": (111, 111),
    "DUAL4": (75, 75),
    "DUALC1": (9, 9),
    "DUALC2": (7, 7),
    "DUALC5": (8, 8),
    "DUALC8": (8, 8),
    "FLETCHER": (4, 0),
    "FLT": (2, 2),
    "GOULDQP1": (56, 56),
    "HADAMARD": (84000, 8000),
    "HIMMELBC": (4, 4),
    "HIMMELBD": (4, 4),
    "HIMMELBE": (5, 5),
    "HS107": (34, 34),
    "HS111": (14, 14),
    "HS111LNP": (14, 14),
    "HS112": (14, 14),
    "HS114": (11, 11),
    "HS119": (53, 53),
    "HS32": (3, 3),
    "HS39": (6, 6),
    "HS40": (7, 7),
    "HS41": (4, 4),
    "HS42": (3, 3),
    "HS46": (6, 6),
    "HS47": (8, 8),
    "HS48": (8, 8),
    "HS49": (6, 6),
    "HS50": (9, 9),
    "HS51": (7, 7),
    "HS52": (7, 7),
    "HS53": (7, 7),
    "HS54": (2, 2),
    "HS55": (15, 15),
    "HS56": (10, 10),
    "HS60": (3, 3),
    "HS61": (4, 2),
    "HS62": (3, 3),
    "HS63": (6, 6),
    "HS68": (4, 4),
    "HS69": (4, 4),
    "HS71": (4, 4),
    "HS73": (4, 4),
    "HS77": (6, 6),
    "HS78": (11, 11),
    "HS79": (8, 8),
    "HS8": (4, 4),
    "HS80": (11, 11),
    "HS81": (11, 11),
    "HS87": (15, 15),
    "LSNNODOC": (10, 10),
    "MARATOS": (2, 2),
    "MSS1": (378, 378),
    "NASH": (157, 157),
    "ODFITS": (15, 15),
    "OPTCNTRL": (70, 70),
    "ORTHREGB": (72, 72),
    "PORTFL1": (12, 12),
    "PORTFL2": (12, 12),
    "PORTFL3": (12, 12),
    "PORTFL4": (12, 12),
    "PORTFL6": (12, 12),
    "S316-322": (2, 0),
    "TAME": (2, 2),
    "TENBARS1": (62, 0),
    "TENBARS2": (62, 0),
    "TENBARS3": (62, 0),
    "TENBARS4": (62, 0),
    "TRO11X3": (9000, 446),
    "TRO3X3": (360, 62),
    "TRO4X4": (1512, 178),
    "TRO5X5": (4320, 298),
    "TRO6X2": (900, 114),
}

EXPECTED_JACOBIAN_INEQ_NNZ: dict[str, tuple[int, int]] = {
    "AVGASA": (30, 30),
    "AVGASB": (30, 30),
    "BATCH": (130, 130),
    "BIGGSC4": (16, 16),
    "BURKEHAN": (1, 1),
    "CANTILVR": (5, 5),
    "CB2": (9, 7),
    "CB3": (9, 7),
    "CHACONN1": (9, 9),
    "CHACONN2": (9, 7),
    "CLEUVEN7": (35948, 35948),
    "CSFI1": (4, 4),
    "CSFI2": (4, 4),
    "DEGENQP": (360, 360),
    "DEGENQPC": (58800, 58800),
    "DEMBO7": (68, 68),
    "DEMYMALO": (9, 9),
    "DIPIGRI": (19, 17),
    "DUALC1": (1926, 1926),
    "DUALC2": (1596, 1596),
    "DUALC5": (2216, 2216),
    "DUALC8": (4016, 4016),
    "ELATTAR": (708, 706),
    "FLETCHER": (6, 6),
    "GIGOMEZ1": (9, 9),
    "GIGOMEZ2": (9, 7),
    "GIGOMEZ3": (9, 7),
    "GOFFIN": (2550, 2550),
    "HADAMARD": (1600, 1600),
    "HAIFAM": (12750, 711),
    "HAIFAS": (31, 18),
    "HATFLDH": (16, 16),
    "HS100": (19, 17),
    "HS101": (31, 31),
    "HS102": (31, 31),
    "HS103": (31, 31),
    "HS104": (17, 17),
    "HS105": (2, 2),
    "HS106": (17, 17),
    "HS108": (39, 19),
    "HS113": (32, 29),
    "HS114": (20, 20),
    "HS116": (43, 43),
    "HS117": (62, 62),
    "HS15": (4, 4),
    "HS16": (4, 4),
    "HS17": (4, 4),
    "HS18": (4, 4),
    "HS19": (4, 4),
    "HS20": (6, 6),
    "HS21": (2, 2),
    "HS21MOD": (2, 2),
    "HS22": (4, 4),
    "HS23": (10, 10),
    "HS24": (6, 6),
    "HS268": (25, 25),
    "HS30": (2, 2),
    "HS31": (2, 2),
    "HS32": (3, 3),
    "HS33": (6, 2),
    "HS34": (4, 4),
    "HS35": (3, 3),
    "HS35I": (3, 3),
    "HS35MOD": (3, 3),
    "HS36": (3, 3),
    "HS37": (6, 6),
    "HS43": (12, 9),
    "HS44": (12, 12),
    "HS44NEW": (12, 12),
    "HS57": (2, 2),
    "HS64": (3, 3),
    "HS65": (3, 2),
    "HS66": (4, 4),
    "HS71": (4, 4),
    "HS72": (8, 8),
    "HS73": (8, 8),
    "HS76": (10, 10),
    "HS76I": (10, 10),
    "HS83": (13, 13),
    "HS86": (37, 37),
    "HS88": (2, 2),
    "HS89": (3, 3),
    "HS90": (4, 4),
    "HS91": (5, 5),
    "HS92": (6, 6),
    "HS93": (12, 12),
    "HS95": (20, 20),
    "HS96": (20, 20),
    "HS97": (20, 20),
    "HS98": (20, 20),
    "LOOTSMA": (6, 6),
    "LSQFIT": (2, 2),
    "MADSEN": (14, 14),
    "MAKELA1": (6, 6),
    "MAKELA2": (9, 8),
    "MAKELA3": (40, 40),
    "MAKELA4": (80, 80),
    "MATRIX2": (6, 6),
    "MINMAXBD": (100, 100),
    "MINMAXRB": (10, 10),
    "OET1": (3006, 3004),
    "OET2": (3006, 2004),
    "OET3": (4008, 4004),
    "OET4": (4008, 3004),
    "OET5": (5010, 5010),
    "OET6": (5010, 3006),
    "OET7": (7014, 4008),
    "OPTPRLOC": (170, 170),
    "PENTAGON": (30, 27),
    "POLAK1": (6, 6),
    "POLAK2": (22, 22),
    "POLAK3": (120, 120),
    "POLAK4": (9, 9),
    "POLAK5": (6, 6),
    "POLAK6": (20, 20),
    "S365": (22, 16),
    "S365MOD": (22, 16),
    "SIMPLLPA": (4, 4),
    "SIMPLLPB": (6, 6),
    "SIPOW1": (4000, 3996),
    "SIPOW2": (3000, 2996),
    "TENBARS1": (2, 2),
    "TENBARS4": (2, 2),
    "TRO11X3": (90, 90),
    "TRO3X3": (18, 18),
    "TRO4X4": (39, 39),
    "TRO5X5": (68, 68),
    "TRO6X2": (25, 25),
    "ZECEVIC2": (4, 4),
    "ZECEVIC3": (4, 4),
    "ZECEVIC4": (4, 4),
}


def _discover_fixtures(kind: str) -> list[str]:
    """Discover fixture problem names by globbing the fixture directory."""
    subdir = FIXTURE_DIR / kind
    if not subdir.exists():
        return []
    return sorted(p.stem for p in subdir.glob("*.npz"))


_unconstrained_by_name = {
    p.name: p for p in sif2jax.unconstrained_minimisation_problems
}
_constrained_by_name = {p.name: p for p in sif2jax.constrained_minimisation_problems}


def _assert_superset(
    detected: np.ndarray,
    ground_truth: np.ndarray,
    name: str,
    kind: str,
):
    """Assert detected is a superset of ground truth, warn on extras.

    Args:
        detected: Dense boolean array from asdex detection.
        ground_truth: Dense boolean array from ground truth.
        name: Problem name for error messages.
        kind: Pattern kind (hessian, jacobian_eq, jacobian_ineq).
    """
    # Hard assertion: no false negatives.
    missed = ground_truth & ~detected
    assert not missed.any(), (
        f"{kind}/{name}: asdex missed {missed.sum()} nonzeros (false negatives)!\n"
        f"  Ground truth nnz: {ground_truth.sum()}\n"
        f"  Detected nnz: {detected.sum()}\n"
        f"  Missed: {missed.sum()}"
    )

    # Warn on conservativeness (extra nonzeros).
    extras = detected & ~ground_truth
    if extras.any():
        total = detected.shape[0] * detected.shape[1]
        warnings.warn(
            f"{kind}/{name}: asdex detected {extras.sum()} extra nonzeros "
            f"(detected={detected.sum()}, truth={ground_truth.sum()}, "
            f"density: {detected.sum() / total:.1%} vs {ground_truth.sum() / total:.1%})",
            stacklevel=2,
        )


def _check_baseline(
    detected_nnz: int,
    name: str,
    kind: str,
    expected_nnz_dict: dict[str, tuple[int, int]],
):
    """Fail on nnz regressions or improvements so baselines stay current."""
    expected = expected_nnz_dict.get(name)
    if expected is None:
        return
    expected_nnz, _ = expected
    key = f"{kind}/{name}"
    if detected_nnz > expected_nnz:
        pytest.fail(
            f"{key}: REGRESSION — detected nnz increased "
            f"from {expected_nnz} to {detected_nnz}"
        )
    if detected_nnz < expected_nnz:
        pytest.fail(
            f"{key}: IMPROVEMENT — detected nnz decreased "
            f"from {expected_nnz} to {detected_nnz}. "
            f"Update the baseline to ({detected_nnz}, {expected[1]})."
        )


def _objective(problem):
    """Build a scalar function from an unconstrained problem's objective."""
    return lambda y: problem.objective(y, problem.args)


def _eq_constraint(problem):
    """Build a vector function from a constrained problem's equality constraints."""
    return lambda y: problem.constraint(y)[0]


def _ineq_constraint(problem):
    """Build a vector function from a constrained problem's inequality constraints."""
    return lambda y: problem.constraint(y)[1]


def _run_sparsity_test(
    name: str,
    kind: str,
    problems_by_name: dict,
    make_fn,
    detect_fn,
    expected_nnz_dict: dict[str, tuple[int, int]],
):
    """Shared test body for CUTEst sparsity tests.

    Args:
        name: CUTEst problem name.
        kind: Fixture subdirectory (hessian, jacobian_eq, jacobian_ineq).
        problems_by_name: Dict mapping problem names to sif2jax problems.
        make_fn: Callable(problem) -> function under test.
        detect_fn: asdex.hessian_sparsity or asdex.jacobian_sparsity.
        expected_nnz_dict: Baseline dict for this kind.
    """
    if name not in problems_by_name:
        pytest.skip(f"{name} not available in sif2jax")
    problem = problems_by_name[name]

    gt = asdex.SparsityPattern.load(FIXTURE_DIR / kind / f"{name}.npz")

    if problem.y0.size != gt.shape[-1]:
        pytest.skip(
            f"{name}: shape mismatch (sif2jax n={problem.y0.size}, fixture n={gt.shape[-1]})"
        )

    try:
        detected = detect_fn(make_fn(problem), problem.y0.shape)
    except NotImplementedError as exc:
        pytest.xfail(f"asdex internal error for {name}: {exc}")

    _assert_superset(
        detected.todense().astype(bool),
        gt.todense().astype(bool),
        name,
        kind,
    )
    _check_baseline(detected.nnz, name, kind, expected_nnz_dict)


# Hessian tests

hessian_names = _discover_fixtures("hessian")


@pytest.mark.slow
@pytest.mark.cutest
@pytest.mark.hessian
@pytest.mark.parametrize("name", hessian_names)
def test_hessian_sparsity_cutest(name):
    """Compare asdex Hessian sparsity against CUTEst ground truth."""
    _run_sparsity_test(
        name,
        "hessian",
        _unconstrained_by_name,
        _objective,
        asdex.hessian_sparsity,
        EXPECTED_HESSIAN_NNZ,
    )


# Jacobian equality constraint tests

jacobian_eq_names = _discover_fixtures("jacobian_eq")


@pytest.mark.slow
@pytest.mark.cutest
@pytest.mark.jacobian
@pytest.mark.parametrize("name", jacobian_eq_names)
def test_jacobian_eq_sparsity_cutest(name):
    """Compare asdex Jacobian sparsity for equality constraints against ground truth."""
    _run_sparsity_test(
        name,
        "jacobian_eq",
        _constrained_by_name,
        _eq_constraint,
        asdex.jacobian_sparsity,
        EXPECTED_JACOBIAN_EQ_NNZ,
    )


# Jacobian inequality constraint tests

jacobian_ineq_names = _discover_fixtures("jacobian_ineq")


@pytest.mark.slow
@pytest.mark.cutest
@pytest.mark.jacobian
@pytest.mark.parametrize("name", jacobian_ineq_names)
def test_jacobian_ineq_sparsity_cutest(name):
    """Compare asdex Jacobian sparsity for inequality constraints against ground truth."""
    _run_sparsity_test(
        name,
        "jacobian_ineq",
        _constrained_by_name,
        _ineq_constraint,
        asdex.jacobian_sparsity,
        EXPECTED_JACOBIAN_INEQ_NNZ,
    )
