"""
Hyperspectral agent pipeline unit tests.

Tests the data flow from controllers through result aggregation to synthesis,
with focus on custom code result passthrough and preprocessing interactions.

No API key needed — all LLM calls are mocked.

Run with:
    python -m pytest tests/test_hyperspectral_pipeline.py -v
"""

import logging
import json
import os
import tempfile
from copy import deepcopy
from unittest.mock import MagicMock, patch, PropertyMock

import numpy as np
import pytest

from scilink.agents.exp_agents.controllers.hyperspectral_controllers import (
    BuildHolisticSynthesisPromptController,
    RunDynamicAnalysisController,
    RunPreprocessingController,
)
from scilink.tools.hyperspectral_tools import (
    estimate_global_snr,
    get_optimal_analysis_data,
)

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

LOGGER = logging.getLogger("test_hyperspectral")


def _make_cube(h=8, w=8, e=32, snr="high"):
    """Create a synthetic hyperspectral cube with controllable SNR."""
    rng = np.random.RandomState(42)
    # Two Gaussian peaks at different spatial locations
    energy = np.linspace(0, 10, e)
    peak1 = np.exp(-((energy - 3) ** 2) / 0.5)
    peak2 = np.exp(-((energy - 7) ** 2) / 0.5)

    cube = np.zeros((h, w, e))
    cube[: h // 2, :, :] = peak1[None, None, :] * 100
    cube[h // 2 :, :, :] = peak2[None, None, :] * 100

    if snr == "high":
        cube += rng.normal(0, 0.01, cube.shape)
    elif snr == "medium":
        cube += rng.normal(0, 8, cube.shape)
    elif snr == "low":
        cube += rng.normal(0, 50, cube.shape)
    return cube, energy


def _make_custom_meta(n=2):
    """Create sample custom analysis metadata list."""
    meta = []
    for i in range(n):
        meta.append({
            "name": f"Feature_{i}",
            "units": "eV",
            "description": f"Test feature {i}",
            "stats": {"min": float(i), "max": float(i + 10), "mean": float(i + 5)},
        })
    return meta


def _make_iteration_result(
    title="Global_Analysis",
    analysis_text="NMF found 3 components.",
    custom_meta=None,
    parent_reasoning=None,
    depth=0,
    n_images=1,
):
    """Create a mock iteration result_summary dict."""
    images = []
    for j in range(n_images):
        images.append({
            "label": f"Plot_{j}",
            "data": b"\xff\xd8\xff\xe0fake_jpeg_bytes",
        })
    return {
        "iteration_title": title,
        "iteration_analysis_text": analysis_text,
        "analysis_images": images,
        "refinement_decision": {},
        "depth": depth,
        "custom_analysis_metadata_list": custom_meta,
        "parent_refinement_reasoning": parent_reasoning,
    }


def _make_synthesis_state(iteration_results, **overrides):
    """Create a synthesis_state dict for BuildHolisticSynthesisPromptController."""
    state = {
        "all_iteration_results": iteration_results,
        "system_info": {"experiment": {"type": "EELS"}},
        "instruction_prompt": "Analyze this data.",
        "analysis_hints": None,
        "analysis_objective": None,
        "prior_knowledge": [],
        "result_json": None,
        "error_dict": None,
        "skill_name": None,
        "skill_sections": None,
        "auxiliary_plot_bytes": None,
        "auxiliary_label": None,
        "auxiliary_summary": None,
        "auxiliary_mime_type": None,
    }
    state.update(overrides)
    return state


# ===========================================================================
# 1. CUSTOM CODE RESULT PASSTHROUGH
# ===========================================================================


class TestCustomCodePassthrough:
    """Verify that custom code results survive from controller → synthesis."""

    def test_custom_meta_included_in_synthesis_prompt(self):
        """The key bug: custom_analysis_metadata_list must appear in the
        final synthesis prompt when present."""
        meta = _make_custom_meta(2)
        results = [_make_iteration_result(custom_meta=meta)]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "DYNAMIC ANALYSIS FINDINGS" in prompt_text
        assert "Feature_0" in prompt_text
        assert "Feature_1" in prompt_text
        assert "eV" in prompt_text

    def test_custom_meta_none_does_not_crash(self):
        """When no custom code was run, synthesis should proceed normally."""
        results = [_make_iteration_result(custom_meta=None)]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "DYNAMIC ANALYSIS FINDINGS" not in prompt_text
        assert out.get("error_dict") is None

    def test_custom_meta_empty_list_treated_as_absent(self):
        """An empty metadata list should be treated same as None."""
        results = [_make_iteration_result(custom_meta=[])]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "DYNAMIC ANALYSIS FINDINGS" not in prompt_text

    def test_custom_meta_stats_missing_gracefully(self):
        """If stats dict is empty or missing keys, synthesis should not crash."""
        meta = [{"name": "Sparse", "units": "nm", "description": "test", "stats": {}}]
        results = [_make_iteration_result(custom_meta=meta)]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)  # Should not raise

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "Sparse" in prompt_text

    def test_custom_meta_stats_none(self):
        """stats=None should not crash the synthesis."""
        meta = [{"name": "NoStats", "units": "a.u.", "description": "x", "stats": None}]
        results = [_make_iteration_result(custom_meta=meta)]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)
        assert out.get("error_dict") is None

    def test_multiple_iterations_each_with_custom_meta(self):
        """Custom meta from every iteration should appear in synthesis."""
        meta_a = [{"name": "Band_Gap", "units": "eV", "description": "d1", "stats": {"min": 1.0, "max": 3.0, "mean": 2.0}}]
        meta_b = [{"name": "Plasmon_Peak", "units": "eV", "description": "d2", "stats": {"min": 10.0, "max": 20.0, "mean": 15.0}}]
        results = [
            _make_iteration_result(title="Global_Analysis", custom_meta=meta_a),
            _make_iteration_result(title="Zoom_Region_1", custom_meta=meta_b, depth=1),
        ]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "Band_Gap" in prompt_text
        assert "Plasmon_Peak" in prompt_text

    def test_mixed_iterations_some_with_custom_some_without(self):
        """Iterations without custom code should not mask those with it."""
        meta = _make_custom_meta(1)
        results = [
            _make_iteration_result(title="Global_Analysis", custom_meta=None),
            _make_iteration_result(title="Zoom_A", custom_meta=meta, depth=1),
        ]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        # Global should NOT have dynamic findings
        # But Zoom_A should
        assert "Feature_0" in prompt_text
        assert "DYNAMIC ANALYSIS FINDINGS" in prompt_text


# ===========================================================================
# 2. PARENT REFINEMENT REASONING PASSTHROUGH
# ===========================================================================


class TestParentReasoningPassthrough:
    """The synthesis prompt should include why each iteration was created."""

    def test_parent_reasoning_appears_in_prompt(self):
        """parent_refinement_reasoning should show up as 'Target Description'."""
        results = [
            _make_iteration_result(
                title="Zoom_Grain_Boundary",
                parent_reasoning="High-loss region near grain boundary suggests phase segregation",
                depth=1,
            )
        ]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "phase segregation" in prompt_text

    def test_none_reasoning_does_not_crash(self):
        """Root iteration has no parent reasoning — should not crash."""
        results = [_make_iteration_result(parent_reasoning=None)]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)
        assert out.get("error_dict") is None


# ===========================================================================
# 3. IMAGE HANDLING IN SYNTHESIS
# ===========================================================================


class TestImageHandling:
    """Verify images get semantic IDs and are collected correctly."""

    def test_images_get_unique_semantic_ids(self):
        results = [
            _make_iteration_result(title="Global_Analysis", n_images=2),
            _make_iteration_result(title="Zoom_A", n_images=1, depth=1),
        ]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        collected = out["analysis_images"]
        labels = [img["label"] for img in collected]
        # Each label should be prefixed with sanitized iteration title
        assert all("[Global_Analysis]" in l for l in labels[:2])
        assert "[Zoom_A]" in labels[2]
        # All unique
        assert len(labels) == len(set(labels))

    def test_images_without_data_are_skipped(self):
        """Images with no 'data' or 'bytes' key should be silently skipped."""
        results = [_make_iteration_result(n_images=0)]
        results[0]["analysis_images"] = [{"label": "Ghost", "data": None}]
        state = _make_synthesis_state(results)

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        assert len(out["analysis_images"]) == 0

    def test_no_results_sets_error(self):
        """Empty iteration results should produce an error, not crash."""
        state = _make_synthesis_state([])

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        assert out.get("error_dict") is not None
        assert "No iteration results" in out["error_dict"]["error"]


# ===========================================================================
# 4. RunDynamicAnalysisController
# ===========================================================================


class TestRunDynamicAnalysisController:
    """Test the custom code execution controller in isolation."""

    def _make_state(self, targets=None, cube_snr="high"):
        cube, energy = _make_cube(snr=cube_snr)
        h, w, e = cube.shape
        state = {
            "refinement_decision": {
                "requires_custom_code": bool(targets),
                "targets": targets or [],
            },
            "hspy_data": cube,
            "energy_axis": energy,
            "system_info": {"energy_range": {"start": 0, "end": 10, "units": "eV"}},
            "settings": {"output_dir": tempfile.mkdtemp()},
            "iteration_title": "Test_Iter",
            "analysis_hints": None,
            "analysis_objective": None,
            "analysis_images": [],
        }
        return state

    def test_skips_when_no_custom_targets(self):
        """Controller should be a no-op when there are no custom_code targets."""
        mock_model = MagicMock()
        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, lambda r: ({}, None)
        )
        state = self._make_state(targets=[])
        out = controller.execute(state)
        assert "custom_analysis_metadata_list" not in out
        mock_model.generate_content.assert_not_called()

    def test_skips_non_custom_targets(self):
        """Targets with type != 'custom_code' should be filtered out."""
        mock_model = MagicMock()
        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, lambda r: ({}, None)
        )
        targets = [{"type": "nmf_zoom", "description": "Zoom into region"}]
        state = self._make_state(targets=targets)
        out = controller.execute(state)
        mock_model.generate_content.assert_not_called()

    def test_successful_custom_code_sets_metadata_list(self):
        """Successful execution should populate custom_analysis_metadata_list."""
        h, w = 8, 8
        # Code that returns a valid map
        code = (
            "def analyze_feature(data, energy_axis):\n"
            "    import numpy as np\n"
            "    h, w, e = data.shape\n"
            "    return {\n"
            "        'maps': {'peak_position': np.ones((h, w)) * 3.0},\n"
            "        'units': 'eV',\n"
            "        'description': 'Peak position map'\n"
            "    }\n"
        )
        mock_model = MagicMock()
        mock_model.generate_content.return_value = MagicMock()

        def parse_fn(response):
            return {"code": code}, None

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        # Patch visual QC to always pass
        controller._check_result_visually = MagicMock(return_value=(True, "OK"))

        targets = [{"type": "custom_code", "description": "Map peak position"}]
        state = self._make_state(targets=targets)
        out = controller.execute(state)

        assert "custom_analysis_metadata_list" in out
        meta = out["custom_analysis_metadata_list"]
        assert len(meta) == 1
        assert meta[0]["name"] == "peak_position"
        assert meta[0]["units"] == "eV"
        assert "stats" in meta[0]
        assert meta[0]["stats"]["min"] == pytest.approx(3.0)

    def test_failed_code_marks_failure(self):
        """If all retries fail, dynamic_analysis_failed should be set."""
        mock_model = MagicMock()

        def parse_fn(response):
            return {"code": "def analyze_feature(data, axis): raise ValueError('boom')"}, None

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller.MAX_RETRIES = 1  # Speed up test

        targets = [{"type": "custom_code", "description": "Bad code"}]
        state = self._make_state(targets=targets)
        out = controller.execute(state)

        assert out.get("dynamic_analysis_failed") is True
        assert "custom_analysis_metadata_list" not in out

    def test_qc_rejection_triggers_retry(self):
        """Visual QC failure should trigger code regeneration."""
        call_count = {"n": 0}

        def parse_fn(response):
            call_count["n"] += 1
            code = (
                "def analyze_feature(data, energy_axis):\n"
                "    import numpy as np\n"
                "    h, w, e = data.shape\n"
                "    return {'maps': {'feat': np.ones((h, w))}, 'units': 'a.u.', 'description': 'test'}\n"
            )
            return {"code": code}, None

        mock_model = MagicMock()

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller.MAX_RETRIES = 2

        # First call: reject; Second call: accept
        controller._check_result_visually = MagicMock(
            side_effect=[(False, "Noisy"), (True, "OK")]
        )

        targets = [{"type": "custom_code", "description": "Test"}]
        state = self._make_state(targets=targets)
        out = controller.execute(state)

        # QC was called twice (reject then accept)
        assert controller._check_result_visually.call_count == 2
        assert "custom_analysis_metadata_list" in out

    def test_wrong_shape_map_skipped(self):
        """Maps with wrong spatial dimensions should be skipped."""
        code = (
            "def analyze_feature(data, energy_axis):\n"
            "    import numpy as np\n"
            "    return {\n"
            "        'maps': {'bad': np.ones((3, 3)), 'good': np.ones((8, 8))},\n"
            "        'units': 'a.u.',\n"
            "        'description': 'test'\n"
            "    }\n"
        )
        mock_model = MagicMock()
        parse_fn = lambda r: ({"code": code}, None)

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller._check_result_visually = MagicMock(return_value=(True, "OK"))

        targets = [{"type": "custom_code", "description": "Mixed shapes"}]
        state = self._make_state(targets=targets)
        out = controller.execute(state)

        meta = out.get("custom_analysis_metadata_list", [])
        names = [m["name"] for m in meta]
        assert "good" in names
        assert "bad" not in names

    def test_all_nan_map_skipped(self):
        """Maps that are entirely NaN should be skipped."""
        code = (
            "def analyze_feature(data, energy_axis):\n"
            "    import numpy as np\n"
            "    h, w, e = data.shape\n"
            "    return {\n"
            "        'maps': {'nan_map': np.full((h, w), np.nan), 'ok_map': np.ones((h, w))},\n"
            "        'units': 'a.u.',\n"
            "        'description': 'test'\n"
            "    }\n"
        )
        mock_model = MagicMock()
        parse_fn = lambda r: ({"code": code}, None)

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller._check_result_visually = MagicMock(return_value=(True, "OK"))

        targets = [{"type": "custom_code", "description": "NaN test"}]
        state = self._make_state(targets=targets)
        out = controller.execute(state)

        meta = out.get("custom_analysis_metadata_list", [])
        names = [m["name"] for m in meta]
        assert "ok_map" in names
        assert "nan_map" not in names


# ===========================================================================
# 5. PREPROCESSING → CUSTOM CODE INTERACTION
# ===========================================================================


class TestPreprocessingCustomCodeInteraction:
    """Verify that preprocessing output is what custom code receives."""

    def test_preprocessed_data_flows_to_custom_code(self):
        """Custom code should run on state['hspy_data'], which may be
        modified by preprocessing. Simulate the full iteration flow."""
        cube, energy = _make_cube()
        h, w, e = cube.shape

        # Simulate preprocessing modifying the data
        preprocessed = cube * 0.5  # e.g., baseline subtraction halved values

        state = {
            "hspy_data": preprocessed,
            "energy_axis": energy,
            "system_info": {"energy_range": {"start": 0, "end": 10, "units": "eV"}},
            "refinement_decision": {
                "requires_custom_code": True,
                "targets": [{"type": "custom_code", "description": "Check values"}],
            },
            "settings": {"output_dir": tempfile.mkdtemp()},
            "iteration_title": "Test",
            "analysis_hints": None,
            "analysis_objective": None,
            "analysis_images": [],
        }

        # Code that captures the mean — should see preprocessed (halved) values
        code = (
            "def analyze_feature(data, energy_axis):\n"
            "    import numpy as np\n"
            "    h, w, e = data.shape\n"
            "    mean_map = np.mean(data, axis=2)\n"
            "    return {'maps': {'mean_intensity': mean_map}, 'units': 'counts', 'description': 'test'}\n"
        )
        mock_model = MagicMock()
        parse_fn = lambda r: ({"code": code}, None)

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller._check_result_visually = MagicMock(return_value=(True, "OK"))

        out = controller.execute(state)
        meta = out["custom_analysis_metadata_list"]

        # The mean should reflect the preprocessed (halved) data, not the original.
        # get_optimal_analysis_data may further PCA-denoise, but the key point
        # is that it starts from state["hspy_data"] (the preprocessed version).
        assert meta[0]["stats"]["mean"] > 0  # Sanity check it ran

    def test_preprocessing_mask_carried_to_refinement(self):
        """Preprocessing mask from depth-0 should be available at depth > 0."""
        cube, _ = _make_cube()
        h, w, _ = cube.shape
        mask = np.ones((h, w), dtype=bool)
        mask[0, 0] = False  # One masked pixel

        # Simulate a depth-1 iteration state (preprocessing skipped)
        mock_preprocessor = MagicMock()
        mock_preprocessor._calculate_statistics.return_value = {"mean": 1.0}
        mock_preprocessor._calculate_snr.return_value = (50.0, "High SNR")

        controller = RunPreprocessingController(LOGGER, mock_preprocessor)

        state = {
            "hspy_data": cube,
            "system_info": {},
            "settings": {"run_preprocessing": False},
            "preprocessing_mask": mask,
        }
        out = controller.execute(state)

        # Mask should be preserved, not overwritten
        assert np.array_equal(out["preprocessing_mask"], mask)
        assert out["preprocessing_mask"][0, 0] == False

    def test_preprocessing_skipped_creates_default_mask(self):
        """When preprocessing is skipped and no mask exists, an all-True mask
        should be created."""
        cube, _ = _make_cube()

        mock_preprocessor = MagicMock()
        mock_preprocessor._calculate_statistics.return_value = {"mean": 1.0}
        mock_preprocessor._calculate_snr.return_value = (50.0, "OK")

        controller = RunPreprocessingController(LOGGER, mock_preprocessor)

        state = {
            "hspy_data": cube,
            "system_info": {},
            "settings": {"run_preprocessing": False},
        }
        out = controller.execute(state)

        assert "preprocessing_mask" in out
        assert out["preprocessing_mask"].all()  # all True


# ===========================================================================
# 6. get_optimal_analysis_data & SNR estimation
# ===========================================================================


class TestGetOptimalAnalysisData:
    """Verify that raw data is always returned unchanged with SNR info."""

    def test_always_returns_raw_data(self):
        """Data must never be modified, regardless of noise level."""
        for snr in ("high", "medium", "low"):
            cube, _ = _make_cube(snr=snr)
            data, note = get_optimal_analysis_data(cube)
            assert np.array_equal(data, cube), f"Data modified for snr={snr}"
            assert "Raw Data" in note

    def test_note_contains_snr_estimate(self):
        cube, _ = _make_cube(snr="high")
        _, note = get_optimal_analysis_data(cube)
        assert "SNR" in note

    def test_note_contains_quality_label(self):
        """Note should report a human-readable quality level."""
        with patch(
            "scilink.tools.hyperspectral_tools.estimate_global_snr", return_value=100.0
        ):
            cube, _ = _make_cube(snr="high")
            _, note = get_optimal_analysis_data(cube)
            assert "High" in note

        with patch(
            "scilink.tools.hyperspectral_tools.estimate_global_snr", return_value=2.0
        ):
            cube, _ = _make_cube(snr="low")
            _, note = get_optimal_analysis_data(cube)
            assert "Very Low" in note

    def test_output_shape_preserved(self):
        for snr in ("high", "medium", "low"):
            cube, _ = _make_cube(h=10, w=12, e=50, snr=snr)
            data, _ = get_optimal_analysis_data(cube)
            assert data.shape == cube.shape


class TestSNREstimator:
    """Verify that MAD-based SNR is robust to sharp spectral features."""

    def test_sharp_peak_does_not_collapse_snr(self):
        """A clean cube with a single sharp peak should still report
        high SNR.  The old std-based estimator returned ~1.3 for this."""
        rng = np.random.RandomState(42)
        h, w, e = 8, 8, 128
        energy = np.linspace(0, 10, e)

        # Smooth baseline + one very sharp peak (3 channels wide)
        spectrum = np.ones(e) * 50.0
        peak_center = 64
        spectrum[peak_center - 1 : peak_center + 2] += 200.0

        cube = np.broadcast_to(spectrum, (h, w, e)).copy()
        cube += rng.normal(0, 0.5, cube.shape)  # Light noise

        snr = estimate_global_snr(cube)
        # With MAD, the sharp peak is an outlier and should not drag SNR down.
        # We expect SNR >> 5 (the old estimator gave ~1.3).
        assert snr > 20.0, f"SNR={snr:.1f} — sharp peak collapsed the estimate"

    def test_huge_peak_does_not_inflate_snr_on_noisy_data(self):
        """A massive peak on a noisy baseline should not make the data
        appear cleaner than it is.  SNR should reflect the baseline
        noise, not be inflated by the peak's contribution to the mean."""
        rng = np.random.RandomState(42)
        h, w, e = 8, 8, 128
        baseline_noise = 10.0

        # Noisy baseline + enormous peak
        spectrum = np.ones(e) * 5.0
        spectrum[60:65] = 2000.0
        cube = np.broadcast_to(spectrum, (h, w, e)).copy()
        cube += rng.normal(0, baseline_noise, cube.shape)

        # Same noise, no peak
        spectrum_flat = np.ones(e) * 5.0
        cube_flat = np.broadcast_to(spectrum_flat, (h, w, e)).copy()
        cube_flat += rng.normal(0, baseline_noise, cube_flat.shape)

        snr_with_peak = estimate_global_snr(cube)
        snr_no_peak = estimate_global_snr(cube_flat)

        # Both should report very low SNR; the peak should not inflate it
        assert snr_with_peak < 5.0, f"SNR={snr_with_peak:.1f} — peak inflated estimate"
        assert snr_with_peak / snr_no_peak < 2.0, (
            f"Peak inflated SNR by {snr_with_peak / snr_no_peak:.1f}x"
        )

    def test_genuinely_noisy_data_reports_low_snr(self):
        """Data dominated by noise should still report low SNR."""
        rng = np.random.RandomState(0)
        cube = rng.normal(10, 10, (8, 8, 64))  # SNR ~ 1
        snr = estimate_global_snr(cube)
        assert snr < 10.0, f"SNR={snr:.1f} — noisy data not detected"

    def test_constant_data_returns_perfect_snr(self):
        cube = np.ones((4, 4, 16)) * 42.0
        snr = estimate_global_snr(cube)
        assert snr == 100.0

    def test_zero_signal_returns_zero(self):
        cube = np.zeros((4, 4, 16))
        snr = estimate_global_snr(cube)
        assert snr == 0.0


# ===========================================================================
# 7. RESULT SUMMARY STRUCTURE (integration-style)
# ===========================================================================


class TestResultSummaryStructure:
    """Verify the result_summary dict built in _run_analysis_pipeline has
    the correct keys for the synthesis controller to consume."""

    REQUIRED_KEYS_FOR_SYNTHESIS = {
        "iteration_title",
        "iteration_analysis_text",
        "analysis_images",
        "custom_analysis_metadata_list",
    }

    def test_result_summary_has_all_synthesis_keys(self):
        """result_summary must contain every key that
        BuildHolisticSynthesisPromptController.execute() reads."""
        # Simulate what the agent builds at lines 704-711
        iteration_state = {
            "iteration_title": "Global_Analysis",
            "result_json": {"detailed_analysis": "Found 3 components."},
            "analysis_images": [{"label": "NMF", "data": b"img"}],
            "refinement_decision": {"action": "zoom"},
            "custom_analysis_metadata_list": _make_custom_meta(1),
        }
        result_summary = {
            "iteration_title": iteration_state.get("iteration_title"),
            "iteration_analysis_text": iteration_state.get("result_json", {}).get("detailed_analysis", ""),
            "analysis_images": iteration_state.get("analysis_images", []),
            "refinement_decision": iteration_state.get("refinement_decision", {}),
            "depth": 0,
            "custom_analysis_metadata_list": iteration_state.get("custom_analysis_metadata_list"),
        }

        for key in self.REQUIRED_KEYS_FOR_SYNTHESIS:
            assert key in result_summary, f"Missing key: {key}"

    def test_parent_reasoning_needed_by_synthesis(self):
        """BuildHolisticSynthesisPromptController reads
        'parent_refinement_reasoning' from iter_result. This key
        must be present in result_summary for refinement iterations."""
        # This test documents the secondary bug: parent_refinement_reasoning
        # is read at line 1172 but may not be in result_summary.
        iteration_state = {
            "iteration_title": "Zoom_A",
            "parent_refinement_reasoning": "High-loss region detected",
            "result_json": {"detailed_analysis": ""},
            "analysis_images": [],
            "refinement_decision": {},
            "custom_analysis_metadata_list": None,
        }

        # Current code builds result_summary — check if parent_refinement_reasoning
        # is included (it should be for synthesis to work)
        result_summary = {
            "iteration_title": iteration_state.get("iteration_title"),
            "iteration_analysis_text": iteration_state.get("result_json", {}).get("detailed_analysis", ""),
            "analysis_images": iteration_state.get("analysis_images", []),
            "refinement_decision": iteration_state.get("refinement_decision", {}),
            "depth": 1,
            "custom_analysis_metadata_list": iteration_state.get("custom_analysis_metadata_list"),
            # NOTE: This key is needed but was missing from the actual code
            "parent_refinement_reasoning": iteration_state.get("parent_refinement_reasoning"),
        }

        # Feed to synthesis controller
        results = [result_summary]
        state = _make_synthesis_state(results)
        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "High-loss region detected" in prompt_text


# ===========================================================================
# 8. EDGE CASES
# ===========================================================================


class TestEdgeCases:
    """Miscellaneous edge cases and robustness checks."""

    def test_synthesis_with_error_dict_is_noop(self):
        """If error_dict is already set, controller should return immediately."""
        state = _make_synthesis_state([_make_iteration_result()])
        state["error_dict"] = {"error": "Previous failure"}

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        assert "final_prompt_parts" not in out

    def test_custom_meta_with_per_map_units(self):
        """Units can be a dict mapping feature_name → unit string."""
        # This tests the RunDynamicAnalysisController's unit handling
        code = (
            "def analyze_feature(data, energy_axis):\n"
            "    import numpy as np\n"
            "    h, w, e = data.shape\n"
            "    return {\n"
            "        'maps': {'width': np.ones((h, w)), 'position': np.ones((h, w)) * 5},\n"
            "        'units': {'width': 'meV', 'position': 'eV'},\n"
            "        'description': 'Peak fitting'\n"
            "    }\n"
        )
        mock_model = MagicMock()
        parse_fn = lambda r: ({"code": code}, None)

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller._check_result_visually = MagicMock(return_value=(True, "OK"))

        cube, energy = _make_cube()
        state = {
            "refinement_decision": {
                "requires_custom_code": True,
                "targets": [{"type": "custom_code", "description": "Peak fit"}],
            },
            "hspy_data": cube,
            "energy_axis": energy,
            "system_info": {"energy_range": {"start": 0, "end": 10, "units": "eV"}},
            "settings": {"output_dir": tempfile.mkdtemp()},
            "iteration_title": "Test",
            "analysis_hints": None,
            "analysis_objective": None,
            "analysis_images": [],
        }
        out = controller.execute(state)

        meta = out["custom_analysis_metadata_list"]
        units_by_name = {m["name"]: m["units"] for m in meta}
        assert units_by_name["width"] == "meV"
        assert units_by_name["position"] == "eV"

    def test_synthesis_with_hints_and_objective(self):
        """Hints and objective should appear in the prompt when provided."""
        results = [_make_iteration_result()]
        state = _make_synthesis_state(
            results,
            analysis_hints="Focus on the Ti L-edge",
            analysis_objective="Determine oxidation state variation",
        )

        controller = BuildHolisticSynthesisPromptController(LOGGER)
        out = controller.execute(state)

        prompt_text = "\n".join(
            p for p in out["final_prompt_parts"] if isinstance(p, str)
        )
        assert "Ti L-edge" in prompt_text
        assert "oxidation state" in prompt_text

    def test_dynamic_controller_no_energy_axis_creates_one(self):
        """If energy_axis is missing, controller should create a channel index."""
        code = (
            "def analyze_feature(data, energy_axis):\n"
            "    import numpy as np\n"
            "    h, w, e = data.shape\n"
            "    return {'maps': {'feat': np.ones((h, w))}, 'units': 'a.u.', 'description': 'x'}\n"
        )
        mock_model = MagicMock()
        parse_fn = lambda r: ({"code": code}, None)

        controller = RunDynamicAnalysisController(
            mock_model, LOGGER, None, None, parse_fn
        )
        controller._check_result_visually = MagicMock(return_value=(True, "OK"))

        cube, _ = _make_cube()
        state = {
            "refinement_decision": {
                "requires_custom_code": True,
                "targets": [{"type": "custom_code", "description": "test"}],
            },
            "hspy_data": cube,
            # No energy_axis key
            "system_info": {},
            "settings": {"output_dir": tempfile.mkdtemp()},
            "iteration_title": "Test",
            "analysis_hints": None,
            "analysis_objective": None,
            "analysis_images": [],
        }
        out = controller.execute(state)

        # Should have created an energy axis
        assert "energy_axis" in out
        assert len(out["energy_axis"]) == cube.shape[2]


# ===========================================================================
# Entry point
# ===========================================================================

if __name__ == "__main__":
    pytest.main([__file__, "-v"])
