"""Unit tests for spey_hs3 helper functions.

These tests exercise the pure dictionary-manipulation and domain-merging helpers
without requiring pyhs3 model compilation, so they are fast.
"""

import copy

import numpy as np
import pytest

from spey_hs3 import (
    _augment_poi,
    _inject_signal,
    _split_data_by_dist,
)

# _select_analysis and _merge_domains require pyhs3 objects, so they are tested
# in test_interface.py alongside the rest of the integration suite.


# ---------------------------------------------------------------------------
# _inject_signal
# ---------------------------------------------------------------------------


class TestInjectSignal:
    def test_list_yields_adds_sample(self, minimal_ws):
        result = _inject_signal(minimal_ws, {"dist1": {"sig": [3.0, 5.0]}}, poi_name="mu")
        samples = result["distributions"][0]["samples"]
        names = [s["name"] for s in samples]
        assert "sig" in names

    def test_list_yields_contents(self, minimal_ws):
        result = _inject_signal(minimal_ws, {"dist1": {"sig": [3.0, 5.0]}}, poi_name="mu")
        sig = next(s for s in result["distributions"][0]["samples"] if s["name"] == "sig")
        assert sig["data"]["contents"] == [3.0, 5.0]

    def test_list_yields_zero_errors(self, minimal_ws):
        result = _inject_signal(minimal_ws, {"dist1": {"sig": [3.0, 5.0]}}, poi_name="mu")
        sig = next(s for s in result["distributions"][0]["samples"] if s["name"] == "sig")
        assert sig["data"]["errors"] == [0.0, 0.0]

    def test_dict_yields_contents_and_errors(self, minimal_ws):
        spec = {"contents": [1.0, 2.0], "errors": [0.1, 0.2]}
        result = _inject_signal(minimal_ws, {"dist1": {"sig": spec}}, poi_name="mu")
        sig = next(s for s in result["distributions"][0]["samples"] if s["name"] == "sig")
        assert sig["data"]["contents"] == [1.0, 2.0]
        assert sig["data"]["errors"] == [0.1, 0.2]

    def test_dict_yields_default_errors(self, minimal_ws):
        spec = {"contents": [1.0, 2.0]}
        result = _inject_signal(minimal_ws, {"dist1": {"sig": spec}}, poi_name="mu")
        sig = next(s for s in result["distributions"][0]["samples"] if s["name"] == "sig")
        assert sig["data"]["errors"] == [0.0, 0.0]

    def test_normfactor_modifier_attached(self, minimal_ws):
        result = _inject_signal(minimal_ws, {"dist1": {"sig": [1.0, 2.0]}}, poi_name="mu")
        sig = next(s for s in result["distributions"][0]["samples"] if s["name"] == "sig")
        mods = sig["modifiers"]
        assert len(mods) == 1
        assert mods[0]["type"] == "normfactor"
        assert mods[0]["parameter"] == "mu"
        assert mods[0]["name"] == "mu"

    def test_does_not_mutate_original(self, minimal_ws):
        original_samples = copy.deepcopy(minimal_ws["distributions"][0]["samples"])
        _inject_signal(minimal_ws, {"dist1": {"sig": [1.0, 2.0]}}, poi_name="mu")
        assert minimal_ws["distributions"][0]["samples"] == original_samples

    def test_poi_added_to_analyses_list(self, minimal_ws):
        result = _inject_signal(minimal_ws, {"dist1": {"sig": [1.0, 2.0]}}, poi_name="mu")
        pois = result["analyses"][0]["parameters_of_interest"]
        assert "mu" in pois

    def test_unknown_dist_raises(self, minimal_ws):
        with pytest.raises(ValueError, match="[Nn]o histfactory_dist"):
            _inject_signal(minimal_ws, {"nonexistent": {"sig": [1.0]}}, poi_name="mu")

    def test_duplicate_sample_raises(self, minimal_ws):
        with pytest.raises(ValueError, match="already exists"):
            _inject_signal(minimal_ws, {"dist1": {"background": [1.0, 2.0]}}, poi_name="mu")

    def test_mismatched_contents_errors_raises(self, minimal_ws):
        spec = {"contents": [1.0, 2.0], "errors": [0.1]}
        with pytest.raises(ValueError, match="same length"):
            _inject_signal(minimal_ws, {"dist1": {"sig": spec}}, poi_name="mu")

    def test_non_histfactory_dist_ignored(self):
        ws = {
            "analyses": [
                {"name": "a", "likelihood": "l", "domains": [], "parameters_of_interest": []}
            ],
            "distributions": [{"name": "other", "type": "not_histfactory", "samples": []}],
            "likelihoods": [],
            "data": [],
            "domains": [],
            "parameter_points": [],
            "metadata": {"hs3_version": "0.2"},
            "misc": {},
        }
        with pytest.raises(ValueError, match="No histfactory_dist"):
            _inject_signal(ws, {"other": {"sig": [1.0]}}, poi_name="mu")

    def test_multiple_samples_injected(self, minimal_ws):
        signal_yields = {"dist1": {"sig1": [1.0, 2.0], "sig2": [3.0, 4.0]}}
        result = _inject_signal(minimal_ws, signal_yields, poi_name="mu")
        names = [s["name"] for s in result["distributions"][0]["samples"]]
        assert "sig1" in names
        assert "sig2" in names


# ---------------------------------------------------------------------------
# _augment_poi
# ---------------------------------------------------------------------------


class TestAugmentPoi:
    def test_adds_poi_to_analyses_list(self):
        data = {"analyses": [{"name": "a", "parameters_of_interest": []}]}
        _augment_poi(data, "mu")
        assert "mu" in data["analyses"][0]["parameters_of_interest"]

    def test_adds_poi_to_analysis_dict(self):
        data = {"analysis": {"parameters_of_interest": []}}
        _augment_poi(data, "mu")
        assert "mu" in data["analysis"]["parameters_of_interest"]

    def test_no_duplicate_in_analyses_list(self):
        data = {"analyses": [{"name": "a", "parameters_of_interest": ["mu"]}]}
        _augment_poi(data, "mu")
        assert data["analyses"][0]["parameters_of_interest"].count("mu") == 1

    def test_no_duplicate_in_analysis_dict(self):
        data = {"analysis": {"parameters_of_interest": ["mu"]}}
        _augment_poi(data, "mu")
        assert data["analysis"]["parameters_of_interest"].count("mu") == 1

    def test_multiple_analyses_all_updated(self):
        data = {
            "analyses": [
                {"name": "a", "parameters_of_interest": []},
                {"name": "b", "parameters_of_interest": []},
            ]
        }
        _augment_poi(data, "mu")
        for analysis in data["analyses"]:
            assert "mu" in analysis["parameters_of_interest"]

    def test_no_analyses_key_is_noop(self):
        data = {"distributions": []}
        _augment_poi(data, "mu")  # should not raise
        assert "analyses" not in data

    def test_non_dict_analysis_entries_skipped(self):
        data = {"analyses": ["not_a_dict"]}
        _augment_poi(data, "mu")  # should not raise


# ---------------------------------------------------------------------------
# _split_data_by_dist
# ---------------------------------------------------------------------------


class TestSplitDataByDist:
    def test_single_dist_two_bins(self):
        flat = np.array([10.0, 20.0])
        dist_par_map = {"d1": ["d1_observed", "mu"]}
        obs_ref = {"d1": np.zeros(2)}
        result = _split_data_by_dist(flat, ["d1"], dist_par_map, obs_ref)
        np.testing.assert_array_equal(result["d1"], [10.0, 20.0])

    def test_two_dists_split_correctly(self):
        flat = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        dist_par_map = {
            "d1": ["d1_observed", "mu"],
            "d2": ["d2_observed", "mu"],
        }
        obs_ref = {"d1": np.zeros(2), "d2": np.zeros(3)}
        result = _split_data_by_dist(flat, ["d1", "d2"], dist_par_map, obs_ref)
        np.testing.assert_array_equal(result["d1"], [1.0, 2.0])
        np.testing.assert_array_equal(result["d2"], [3.0, 4.0, 5.0])

    def test_dist_without_observed_param_skipped(self):
        flat = np.array([7.0, 8.0])
        dist_par_map = {
            "d1": ["mu"],   # no _observed param
            "d2": ["d2_observed"],
        }
        obs_ref = {"d2": np.zeros(2)}
        result = _split_data_by_dist(flat, ["d1", "d2"], dist_par_map, obs_ref)
        assert "d1" not in result
        np.testing.assert_array_equal(result["d2"], [7.0, 8.0])

    def test_unknown_obs_ref_defaults_to_one_bin(self):
        flat = np.array([42.0])
        dist_par_map = {"d1": ["d1_observed"]}
        obs_ref = {}  # no reference data
        result = _split_data_by_dist(flat, ["d1"], dist_par_map, obs_ref)
        np.testing.assert_array_equal(result["d1"], [42.0])
