"""Tests comparing JAX function output against MIR interpreter

This module validates that the JAX-compiled models produce the same
results as the reference MIR interpreter.
"""

import numpy as np
import pytest
from conftest import (
    INTEGRATION_MODELS,
    INTEGRATION_PATH,
    CompiledModel,
    assert_allclose,
)
from test_all_models import HANGING_MODELS


class TestJaxVsInterpreter:
    """Compare JAX function output against MIR interpreter for all models"""

    @pytest.mark.parametrize("model_name,model_path", INTEGRATION_MODELS)
    def test_model_compiles(self, compile_model, model_name, model_path):
        """Model compiles to JAX without error"""
        model = compile_model(INTEGRATION_PATH / model_path)
        assert model.module is not None
        assert model.jax_fn is not None
        assert model.name, f"{model_name} has no module name"

    # Simple models that work with the current JAX translator
    SIMPLE_MODELS = ["resistor", "diode", "isrc", "vccs", "cccs"]

    @pytest.mark.parametrize(
        "model_name,model_path",
        [m for m in INTEGRATION_MODELS if m[0] in ["resistor", "diode", "isrc", "vccs", "cccs"]],
    )
    def test_simple_model_produces_valid_output(self, compile_model, model_name, model_path):
        """Simple JAX function produces non-NaN outputs"""
        model = compile_model(INTEGRATION_PATH / model_path)
        inputs = model.build_default_inputs()

        residuals, jacobian = model.jax_fn(inputs)

        # Check residuals are valid
        assert residuals is not None, f"{model_name} returned None residuals"
        for node, res in residuals.items():
            resist = float(res["resist"])
            react = float(res["react"])
            assert not np.isnan(resist), f"{model_name} NaN resist at {node}"
            assert not np.isnan(react), f"{model_name} NaN react at {node}"

    # All complex models now work with proper VA defaults
    # Note: HANGING_MODELS excluded due to CI timeout issues
    # Note: bsimcmg excluded - known to have NaN at internal nodes (see test_all_models.py)
    COMPLEX_MODELS = [
        "diode_cmc",
        "ekv",
        "psp102",
        "psp103",
        "juncap",
        "asmhemt",
        "mvsg",
        "bsim3",
        "bsim4",
        "bsim6",
        "bsimbulk",
        "hicum",
        "mextram",
    ]

    @pytest.mark.parametrize(
        "model_name,model_path",
        [
            m
            for m in INTEGRATION_MODELS
            if m[0]
            in [
                "diode_cmc",
                "ekv",
                "psp102",
                "psp103",
                "juncap",
                "asmhemt",
                "mvsg",
                "bsim3",
                "bsim4",
                "bsim6",
                "bsimbulk",
                "hicum",
                "mextram",
            ]
            and m[0] not in HANGING_MODELS
        ],
    )
    def test_complex_model_produces_valid_output(self, compile_model, model_name, model_path):
        """Complex JAX function produces non-NaN resist outputs with VA defaults"""
        model = compile_model(INTEGRATION_PATH / model_path)
        inputs = model.build_default_inputs()

        residuals, jacobian = model.jax_fn(inputs)

        # Check resist residuals are valid (react may have NaN for some models)
        assert residuals is not None, f"{model_name} returned None residuals"
        for node, res in residuals.items():
            resist = float(res["resist"])
            assert not np.isnan(resist), f"{model_name} NaN resist at {node}"


class TestResistorJaxInterpreter:
    """Detailed comparison for resistor model"""

    @pytest.mark.parametrize(
        "voltage,resistance,temperature,tnom,zeta,mfactor",
        [
            (1.0, 1000.0, 300.0, 300.0, 0.0, 1.0),
            (5.0, 1000.0, 300.0, 300.0, 0.0, 1.0),
            (1.0, 470.0, 300.0, 300.0, 0.0, 1.0),
            (1.0, 1000.0, 350.0, 300.0, 1.0, 1.0),
            (1.0, 1000.0, 350.0, 300.0, 2.0, 1.0),
            (1.0, 1000.0, 300.0, 300.0, 0.0, 2.0),
            (1.0, 1000.0, 250.0, 300.0, 1.0, 1.0),
            (0.1, 100.0, 273.15, 300.0, 1.5, 0.5),
        ],
    )
    def test_resistor_residuals_match(
        self, resistor_model: CompiledModel, voltage, resistance, temperature, tnom, zeta, mfactor
    ):
        """Compare JAX vs interpreter residuals for resistor"""
        # Build inputs for JAX
        jax_inputs = [voltage, voltage, resistance, temperature, tnom, zeta, resistance, mfactor]

        # Build params for interpreter
        interp_params = {
            "V(A,B)": voltage,
            "vres": voltage,
            "R": resistance,
            "$temperature": temperature,
            "tnom": tnom,
            "zeta": zeta,
            "res": resistance,
            "mfactor": mfactor,
        }

        # Run both
        jax_residuals, jax_jacobian = resistor_model.jax_fn(jax_inputs)
        interp_residuals, interp_jacobian = resistor_model.module.run_init_eval(interp_params)

        # Compare residual
        jax_I = float(jax_residuals["A"]["resist"])
        interp_I = interp_residuals[0][0]  # First node, resist component

        assert_allclose(
            jax_I,
            interp_I,
            rtol=1e-6,
            atol=1e-15,
            err_msg=f"Resistor residual mismatch: V={voltage}, R={resistance}",
        )

    @pytest.mark.parametrize(
        "voltage,resistance,temperature,tnom,zeta,mfactor",
        [
            (1.0, 1000.0, 300.0, 300.0, 0.0, 1.0),
            (5.0, 100.0, 300.0, 300.0, 0.0, 1.0),
            (1.0, 1000.0, 350.0, 300.0, 1.0, 1.0),
        ],
    )
    def test_resistor_jacobian_match(
        self, resistor_model: CompiledModel, voltage, resistance, temperature, tnom, zeta, mfactor
    ):
        """Compare JAX vs interpreter jacobian for resistor"""
        jax_inputs = [voltage, voltage, resistance, temperature, tnom, zeta, resistance, mfactor]

        interp_params = {
            "V(A,B)": voltage,
            "vres": voltage,
            "R": resistance,
            "$temperature": temperature,
            "tnom": tnom,
            "zeta": zeta,
            "res": resistance,
            "mfactor": mfactor,
        }

        jax_residuals, jax_jacobian = resistor_model.jax_fn(jax_inputs)
        interp_residuals, interp_jacobian = resistor_model.module.run_init_eval(interp_params)

        # Compare jacobian entry (0,0)
        jax_jac_00 = float(jax_jacobian[("A", "A")]["resist"])
        interp_jac_00 = interp_jacobian[0][2]  # row=0, col=0, resist component

        assert_allclose(
            jax_jac_00,
            interp_jac_00,
            rtol=1e-5,
            atol=1e-12,
            err_msg=f"Resistor jacobian mismatch: R={resistance}",
        )


class TestModelNodeCounts:
    """Verify models have reasonable node counts"""

    @pytest.mark.parametrize("model_name,model_path", INTEGRATION_MODELS)
    def test_has_nodes(self, compile_model, model_name, model_path):
        """Model has at least 2 nodes"""
        model = compile_model(INTEGRATION_PATH / model_path)
        assert len(model.nodes) >= 2, f"{model_name} should have at least 2 nodes"

    @pytest.mark.parametrize(
        "model_name,model_path",
        [m for m in INTEGRATION_MODELS if m[0] in ["resistor", "isrc", "juncap"]],
    )
    def test_two_terminal_devices(self, compile_model, model_name, model_path):
        """Two-terminal devices have 2 nodes"""
        model = compile_model(INTEGRATION_PATH / model_path)
        assert len(model.nodes) == 2, f"{model_name} should be a two-terminal device"

    @pytest.mark.parametrize(
        "model_name,model_path", [m for m in INTEGRATION_MODELS if m[0] in ["vccs", "cccs"]]
    )
    def test_four_terminal_devices(self, compile_model, model_name, model_path):
        """Controlled sources have at least 4 terminals"""
        model = compile_model(INTEGRATION_PATH / model_path)
        assert len(model.nodes) >= 4, f"{model_name} should have at least 4 terminals"


class TestModelComplexity:
    """Test that complex models compile and produce outputs"""

    @pytest.mark.parametrize(
        "model_name,model_path",
        [
            m
            for m in INTEGRATION_MODELS
            if m[0] in ("bsim4", "psp103", "hisim2", "hicum", "mextram")
        ],
    )
    def test_complex_model_compiles(self, compile_model, model_name, model_path):
        """Complex model compiles to JAX"""
        model = compile_model(INTEGRATION_PATH / model_path)
        assert model.jax_fn is not None
        assert model.module.num_jacobian > 0

    @pytest.mark.parametrize(
        "model_name,model_path",
        [
            m
            for m in INTEGRATION_MODELS
            if m[0] in ("psp103", "bsim4", "hicum", "mextram")  # These now work (hisim2 hangs)
            and m[0] not in HANGING_MODELS
        ],
    )
    def test_working_complex_model_outputs(self, compile_model, model_name, model_path):
        """Working complex model produces finite outputs"""
        model = compile_model(INTEGRATION_PATH / model_path)
        inputs = model.build_default_inputs()

        residuals, jacobian = model.jax_fn(inputs)

        # Check at least one output is finite
        has_finite = False
        for node, res in residuals.items():
            if np.isfinite(float(res["resist"])):
                has_finite = True
                break

        assert has_finite, f"{model_name} produced no finite outputs"

    # Note: bsim4, hisim2, hicum, mextram all produce finite outputs now
    # and are tested in test_working_complex_model_outputs above
