"""Integration tests for HS3Interface and pyhs3-dependent helpers.

These tests compile pyhs3 models and are therefore slower than the pure
helper tests in test_helpers.py.  Mark them explicitly so they can be
excluded during quick iteration:

    pytest -m "not integration"
"""

import copy

import numpy as np
import pytest
import pyhs3
from pyhs3.domains import ProductDomain

import spey
from spey.utils import ExpectationType

from spey_hs3 import HS3Interface, _merge_domains
from tests.conftest import UNCORR_FULL_WORKSPACE, UNCORR_BG_ONLY_WORKSPACE

pytestmark = pytest.mark.integration


# ---------------------------------------------------------------------------
# Analysis selection (via HS3Interface using WorkspaceInterpreter internally)
# ---------------------------------------------------------------------------


class TestAnalysisSelection:
    """Analysis selection is now handled via WorkspaceInterpreter inside HS3Interface."""

    def test_default_analysis_selected(self):
        iface = HS3Interface(
            hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
            analysis_name="simPdf_asimovData",
        )
        assert iface._poi_name == "mu"

    def test_named_analysis_selected(self):
        iface = HS3Interface(
            hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
            analysis_name="simPdf_obsData",
        )
        assert iface._poi_name == "mu"

    def test_raises_for_unknown_analysis(self):
        with pytest.raises((ValueError, AssertionError)):
            HS3Interface(
                hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
                analysis_name="does_not_exist",
            )


# ---------------------------------------------------------------------------
# _merge_domains
# ---------------------------------------------------------------------------


class TestMergeDomains:
    @pytest.fixture(scope="class")
    def ws(self):
        return pyhs3.Workspace(**copy.deepcopy(UNCORR_FULL_WORKSPACE))

    def test_returns_product_domain(self, ws):
        merged = _merge_domains(ws, ["simPdf_asimovData_nuisance_parameters"])
        assert isinstance(merged, ProductDomain)

    def test_merged_axes_contain_nuisance_params(self, ws):
        merged = _merge_domains(
            ws,
            [
                "simPdf_asimovData_nuisance_parameters",
                "simPdf_asimovData_parameters_of_interest",
            ],
        )
        axis_names = {a.name for a in merged.axes}
        assert "uncorr_bkguncrt_0" in axis_names
        assert "uncorr_bkguncrt_1" in axis_names
        assert "mu" in axis_names

    def test_later_domain_overrides_axis(self, ws):
        # Both domains contain "mu"; the second should win (same value here).
        # We just check that there is only one "mu" axis in the result.
        merged = _merge_domains(
            ws,
            [
                "simPdf_asimovData_parameters_of_interest",
                "simPdf_obsData_parameters_of_interest",
            ],
        )
        mu_axes = [a for a in merged.axes if a.name == "mu"]
        assert len(mu_axes) == 1

    def test_empty_domain_names_gives_empty_axes(self, ws):
        merged = _merge_domains(ws, [])
        assert list(merged.axes) == []


# ---------------------------------------------------------------------------
# HS3Interface — instantiation
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def iface_full():
    """HS3Interface compiled from the full workspace (has signal sample)."""
    return HS3Interface(
        hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
        analysis_name="simPdf_asimovData",
    )


@pytest.fixture(scope="module")
def iface_obs():
    """HS3Interface using the observed-data analysis."""
    return HS3Interface(
        hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
        analysis_name="simPdf_obsData",
    )


@pytest.fixture(scope="module")
def iface_with_signal():
    """HS3Interface with externally injected signal on a background-only workspace."""
    return HS3Interface(
        hs3_dict=copy.deepcopy(UNCORR_BG_ONLY_WORKSPACE),
        signal_yields={"model_singlechannel": {"injected_signal": [5.0, 4.0]}},
        poi_name="mu",
        analysis_name="simPdf_asimovData",
    )


@pytest.fixture(scope="module")
def iface_zero_signal():
    """HS3Interface with all-zero signal (is_alive should return False)."""
    return HS3Interface(
        hs3_dict=copy.deepcopy(UNCORR_BG_ONLY_WORKSPACE),
        signal_yields={"model_singlechannel": {"injected_signal": [0.0, 0.0]}},
        poi_name="mu",
        analysis_name="simPdf_asimovData",
    )


class TestHS3InterfaceInstantiation:
    def test_name_attribute(self, iface_full):
        assert iface_full.name == "hs3"

    def test_poi_name_default(self, iface_full):
        assert iface_full._poi_name == "mu"

    def test_explicit_poi_name(self):
        iface = HS3Interface(
            hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
            poi_name="mu",
            analysis_name="simPdf_asimovData",
        )
        assert iface._poi_name == "mu"

    def test_free_params_poi_is_first(self, iface_full):
        assert iface_full._free_params[0] == "mu"

    def test_free_params_excludes_const(self, iface_full):
        # Lumi is const=True in parameter_points
        assert "Lumi" not in iface_full._free_params

    def test_const_params_contains_lumi(self, iface_full):
        assert "Lumi" in iface_full._const_params
        assert iface_full._const_params["Lumi"] == pytest.approx(1.0)

    def test_observed_data_shape(self, iface_full):
        obs = iface_full._observed_data.get("model_singlechannel")
        assert obs is not None
        assert len(obs) == 2  # 2-bin model

    def test_no_poi_raises(self):
        ws_dict = copy.deepcopy(UNCORR_BG_ONLY_WORKSPACE)
        with pytest.raises(ValueError, match="parameters_of_interest"):
            HS3Interface(hs3_dict=ws_dict, analysis_name="simPdf_asimovData")

    def test_registered_as_spey_backend(self):
        # spey.get_backend returns a factory that wraps the backend in a
        # StatisticalModel; verify the underlying backend is an HS3Interface.
        backend_factory = spey.get_backend("hs3")
        stat_model = backend_factory(
            hs3_dict=copy.deepcopy(UNCORR_FULL_WORKSPACE),
            analysis_name="simPdf_asimovData",
        )
        assert isinstance(stat_model._backend, HS3Interface)

    def test_analysis_selection(self, iface_obs):
        assert iface_obs._poi_name == "mu"


# ---------------------------------------------------------------------------
# is_alive
# ---------------------------------------------------------------------------


class TestIsAlive:
    def test_true_for_workspace_with_builtin_signal(self, iface_full):
        # No explicit signal_yields; built-in signal in workspace -> always True
        assert iface_full.is_alive is True

    def test_true_when_injected_signal_nonzero(self, iface_with_signal):
        assert iface_with_signal.is_alive is True

    def test_false_when_all_zero_signal(self, iface_zero_signal):
        assert iface_zero_signal.is_alive is False


# ---------------------------------------------------------------------------
# config()
# ---------------------------------------------------------------------------


class TestConfig:
    def test_poi_index_is_zero(self, iface_full):
        cfg = iface_full.config()
        assert cfg.poi_index == 0

    def test_number_of_params(self, iface_full):
        cfg = iface_full.config()
        # mu + uncorr_bkguncrt_0 + uncorr_bkguncrt_1 = 3 free params
        assert len(cfg.suggested_init) == 3
        assert len(cfg.suggested_bounds) == 3
        assert len(cfg.parameter_names) == 3

    def test_poi_is_first_parameter_name(self, iface_full):
        cfg = iface_full.config()
        assert cfg.parameter_names[0] == "mu"

    def test_default_poi_bounds(self, iface_full):
        cfg = iface_full.config()
        lo, hi = cfg.suggested_bounds[0]
        assert lo < 0  # allow_negative_signal=True by default
        assert hi == pytest.approx(10.0)

    def test_no_negative_signal_bound(self, iface_full):
        cfg = iface_full.config(allow_negative_signal=False)
        lo, _ = cfg.suggested_bounds[0]
        assert lo == pytest.approx(0.0)

    def test_custom_poi_upper_bound(self, iface_full):
        cfg = iface_full.config(poi_upper_bound=5.0)
        _, hi = cfg.suggested_bounds[0]
        assert hi == pytest.approx(5.0)

    def test_minimum_poi_respects_negative(self, iface_full):
        cfg = iface_full.config(allow_negative_signal=True, poi_upper_bound=10.0)
        assert cfg.minimum_poi == pytest.approx(-10.0)

    def test_nuisance_bounds_from_domain(self, iface_full):
        cfg = iface_full.config()
        # uncorr_bkguncrt_0 domain: [0, 1.3]
        nuis_idx = cfg.parameter_names.index("uncorr_bkguncrt_0")
        lo, hi = cfg.suggested_bounds[nuis_idx]
        assert lo == pytest.approx(0.0)
        assert hi == pytest.approx(1.3)


# ---------------------------------------------------------------------------
# get_logpdf_func()
# ---------------------------------------------------------------------------


class TestLogpdf:
    @pytest.fixture(scope="class")
    def logpdf_obs(self, iface_full):
        return iface_full.get_logpdf_func(expected=ExpectationType.observed)

    @pytest.fixture(scope="class")
    def logpdf_apriori(self, iface_full):
        return iface_full.get_logpdf_func(expected=ExpectationType.apriori)

    @pytest.fixture(scope="class")
    def nominal_pars(self, iface_full):
        return np.array(
            [iface_full._init_values.get(p, 1.0) for p in iface_full._free_params]
        )

    def test_returns_callable(self, logpdf_obs):
        assert callable(logpdf_obs)

    def test_logpdf_finite_at_nominal(self, logpdf_obs, nominal_pars):
        val = logpdf_obs(nominal_pars)
        assert np.isfinite(val), f"logpdf={val} is not finite at nominal params"

    def test_logpdf_is_scalar(self, logpdf_obs, nominal_pars):
        val = logpdf_obs(nominal_pars)
        assert np.ndim(val) == 0 or (hasattr(val, "__len__") and len(val) == 1)

    def test_apriori_logpdf_finite(self, logpdf_apriori, nominal_pars):
        val = logpdf_apriori(nominal_pars)
        assert np.isfinite(val)

    def test_logpdf_changes_with_poi(self, logpdf_obs, nominal_pars):
        pars_mu0 = nominal_pars.copy()
        pars_mu0[0] = 0.0
        pars_mu2 = nominal_pars.copy()
        pars_mu2[0] = 2.0
        assert logpdf_obs(pars_mu0) != logpdf_obs(pars_mu2)

    def test_external_data_override(self, iface_full, nominal_pars):
        # Provide flat external data with same number of bins
        ext_data = np.array([55.0, 60.0], dtype=np.float64)
        logpdf_ext = iface_full.get_logpdf_func(
            expected=ExpectationType.observed, data=ext_data
        )
        logpdf_obs = iface_full.get_logpdf_func(expected=ExpectationType.observed)
        val_ext = logpdf_ext(nominal_pars)
        val_obs = logpdf_obs(nominal_pars)
        assert np.isfinite(val_ext)
        # Different data should yield different log-likelihood values
        assert val_ext != val_obs


# ---------------------------------------------------------------------------
# expected_data()
# ---------------------------------------------------------------------------


class TestExpectedData:
    @pytest.fixture(scope="class")
    def nominal_pars(self, iface_full):
        return [iface_full._init_values.get(p, 1.0) for p in iface_full._free_params]

    def test_returns_list(self, iface_full, nominal_pars):
        result = iface_full.expected_data(nominal_pars)
        assert isinstance(result, list)

    def test_length_matches_bins(self, iface_full, nominal_pars):
        result = iface_full.expected_data(nominal_pars)
        n_bins = sum(
            len(v) for v in iface_full._observed_data.values()
        )
        assert len(result) == n_bins

    def test_values_are_finite(self, iface_full, nominal_pars):
        result = iface_full.expected_data(nominal_pars)
        assert all(np.isfinite(v) for v in result)

    def test_values_are_positive(self, iface_full, nominal_pars):
        result = iface_full.expected_data(nominal_pars)
        assert all(v > 0 for v in result), f"Non-positive expected counts: {result}"

    def test_mu_zero_reduces_expected(self, iface_full):
        pars_mu1 = np.array(
            [iface_full._init_values.get(p, 1.0) for p in iface_full._free_params]
        )
        pars_mu0 = pars_mu1.copy()
        pars_mu0[0] = 0.0

        exp_mu1 = iface_full.expected_data(pars_mu1.tolist())
        exp_mu0 = iface_full.expected_data(pars_mu0.tolist())
        # With signal present, mu=0 should give lower expected counts than mu=1
        assert sum(exp_mu0) < sum(exp_mu1)
