"""B-scan rendering and cache helpers."""

from __future__ import annotations

from collections import deque

from PyQt6.QtCore import QRectF
import numpy as np
import pyqtgraph as pg

from python_app.models.dataset_model import ResultCollection
from python_app.orchestration.live_processing_config import ProcessingLiveConfig


def _result_tail(
    *,
    result_history: list[ResultCollection],
    history_limit: int,
    floor_collection_id: int,
) -> list[ResultCollection]:
    """Return filtered and de-duplicated result-history tail for B-scan usage."""
    filtered = [
        collection
        for collection in result_history[-history_limit:]
        if int(collection.collection_id) > int(floor_collection_id)
    ]
    unique_reversed_tail: list[ResultCollection] = []
    seen_keys: set[tuple[int, int]] = set()
    for collection in reversed(filtered):
        key = (int(collection.collection_id), int(collection.monotonic_ns))
        if key in seen_keys:
            continue
        seen_keys.add(key)
        unique_reversed_tail.append(collection)

    unique_reversed_tail.reverse()
    return unique_reversed_tail


def build_bscan_signature(
    live_config: ProcessingLiveConfig,
    *,
    subtract_mean_ascan_enabled: bool,
    result_history: list[ResultCollection],
    history_limit: int,
    floor_collection_id: int,
) -> tuple[object, ...]:
    """Build deterministic signature used to detect B-scan cache invalidation."""
    result_tail = _result_tail(
        result_history=result_history,
        history_limit=history_limit,
        floor_collection_id=floor_collection_id,
    )
    return (
        str(live_config.bscan_axis),
        str(live_config.bscan_channel),
        float(live_config.bscan_cut_m),
        float(live_config.bscan_max_depth_m),
        float(live_config.bscan_gain),
        float(live_config.bscan_start_freq_mhz),
        float(live_config.bscan_stop_freq_mhz),
        bool(subtract_mean_ascan_enabled),
        int(floor_collection_id),
        tuple((int(collection.collection_id), int(collection.monotonic_ns), len(collection.blocks)) for collection in result_tail),
    )


def apply_mean_ascan_subtraction(
    history: deque[np.ndarray],
    *,
    enabled: bool,
) -> np.ndarray:
    """Subtract mean A-scan across sweeps when enabled."""
    sweeps = np.vstack(history).astype(np.float32, copy=False)
    if not enabled:
        return sweeps
    mean_trace = np.mean(sweeps, axis=0, dtype=np.float32)
    return sweeps - mean_trace


def rebuild_bscan_history_from_results(
    result_history: list[ResultCollection],
    history_limit: int,
    floor_collection_id: int,
) -> tuple[dict[tuple[int, int], deque[np.ndarray]], dict[tuple[int, int], np.ndarray]]:
    """Rebuild B-scan history and depth axes from processed result payloads."""
    history_by_combo: dict[tuple[int, int], deque[np.ndarray]] = {}
    depth_axis_by_combo: dict[tuple[int, int], np.ndarray] = {}

    result_tail = _result_tail(
        result_history=result_history,
        history_limit=history_limit,
        floor_collection_id=floor_collection_id,
    )

    for collection in result_tail:
        for block in collection.blocks:
            key = (block.combo.input_pos, block.combo.output_pos)
            for payload in block.payloads:
                if payload.kind != 1 or payload.processing_name != "bscan":
                    continue
                if payload.frequency_hz.size == 0 or payload.trace.size == 0:
                    continue
                if payload.frequency_hz.size != payload.trace.size:
                    continue

                depth_axis = np.asarray(payload.frequency_hz, dtype=np.float32)
                amplitudes = np.asarray(np.real(payload.trace), dtype=np.float32)
                if depth_axis.size == 0 or amplitudes.size == 0:
                    continue

                history = history_by_combo.get(key)
                stored_axis = depth_axis_by_combo.get(key)
                if (
                    history is None
                    or stored_axis is None
                    or stored_axis.shape != depth_axis.shape
                    or not np.allclose(stored_axis, depth_axis, rtol=1e-4, atol=1e-6)
                ):
                    history = deque(maxlen=history_limit)
                    history_by_combo[key] = history
                    depth_axis_by_combo[key] = depth_axis.copy()

                history.append(amplitudes.copy())

    return history_by_combo, depth_axis_by_combo


def pick_bscan_display_key(
    history_by_combo: dict[tuple[int, int], deque[np.ndarray]],
    preferred_key: tuple[int, int] | None = None,
) -> tuple[int, int] | None:
    """Choose combo key to display when multiple histories are present."""
    if not history_by_combo:
        return None
    if preferred_key is not None and preferred_key in history_by_combo:
        return preferred_key
    return next(iter(history_by_combo.keys()))


def bscan_lookup_table(axis_mode: str) -> np.ndarray:
    """Build B-scan colormap table for selected axis mode."""
    if axis_mode == "abs":
        return build_lut(["#440154", "#31688e", "#35b779", "#fde725"])
    return build_lut(["#2166ac", "#67a9cf", "#f7f7f7", "#ef8a62", "#b2182b"])


def build_lut(stops: list[str], *, size: int = 256) -> np.ndarray:
    """Interpolate hex color stops into 8-bit RGB LUT array."""
    stop_positions = np.linspace(0.0, 1.0, num=len(stops), dtype=np.float32)
    sample_positions = np.linspace(0.0, 1.0, num=size, dtype=np.float32)
    stop_colors = np.asarray([pg.mkColor(value).getRgb()[:3] for value in stops], dtype=np.float32)

    lut = np.empty((size, 3), dtype=np.uint8)
    for channel in range(3):
        lut[:, channel] = np.interp(sample_positions, stop_positions, stop_colors[:, channel]).astype(np.uint8)
    return lut


def bscan_levels(sweeps: np.ndarray, axis_mode: str) -> tuple[float, float]:
    """Compute image levels for B-scan data based on axis mode."""
    min_value = float(np.min(sweeps))
    max_value = float(np.max(sweeps))
    if axis_mode == "abs":
        if max_value <= min_value:
            return min_value, min_value + 1e-6
        return min_value, max_value

    max_abs = max(abs(min_value), abs(max_value), 1e-6)
    return -max_abs, max_abs


class AppWindowBscanPlotMixin:
    """Renders B-scan heatmaps and maintains B-scan history caches."""

    def _draw_bscan_heatmap(self, _collection: ResultCollection) -> bool:
        """Draw B-scan image rebuilt from processed result history."""
        self._disable_phase_axis()
        self._sync_bscan_history_from_results()
        return self._draw_bscan_heatmap_from_history()

    def _configure_bscan_plot_axes(self) -> None:
        """Apply persistent B-scan plot axis labels and base view settings."""
        plot = self._bscan_plot
        plot_item = plot.getPlotItem()
        plot_item.showAxis("left", show=True)
        plot_item.showAxis("bottom", show=True)
        plot.setLabel("bottom", "Sweep #")
        plot.setLabel("left", "Range", units="m")
        view_box = plot.getViewBox()
        view_box.invertY(False)
        view_box.enableAutoRange(x=False, y=False)

    def _draw_bscan_heatmap_from_history(self) -> bool:
        """Render B-scan heatmap from currently cached history arrays."""
        display_key = self._pick_bscan_display_key()
        if display_key is None:
            return False

        history = self._bscan_history_by_combo.get(display_key)
        depth_axis = self._bscan_depth_axis_by_combo.get(display_key)
        if not history or depth_axis is None:
            return False

        sweeps = apply_mean_ascan_subtraction(
            history,
            enabled=bool(self._bscan_subtract_mean_ascan.isChecked()),
        )
        if sweeps.size == 0:
            return False

        depth_min = float(np.min(depth_axis))
        depth_max = float(np.max(depth_axis))
        depth_span = max(depth_max - depth_min, 1e-6)
        sweep_count = sweeps.shape[0]
        sweep_width = float(max(sweep_count, 1))
        x_min = 0.5
        x_max = x_min + sweep_width

        image_item = pg.ImageItem(axisOrder="row-major")
        image_item.setImage(sweeps.T, autoLevels=False)
        image_item.setRect(QRectF(x_min, depth_min, sweep_width, depth_span))

        axis_mode = self._bscan_axis.currentText()
        image_item.setLookupTable(self._bscan_lookup_table(axis_mode))
        image_item.setLevels(self._bscan_levels(sweeps, axis_mode))

        self._bscan_plot.clear()
        self._configure_bscan_plot_axes()
        self._bscan_plot.addItem(image_item)
        self._bscan_plot.setXRange(x_min, x_max, padding=0.02)
        self._bscan_plot.setYRange(depth_min, depth_max, padding=0.02)
        bscan_channel = "S21"
        self._bscan_plot.setTitle(
            f"B-scan {bscan_channel} in{display_key[0]}/out{display_key[1]} | sweeps={sweep_count}"
        )
        return True

    def _sync_bscan_history_from_results(self) -> None:
        """Rebuild B-scan history cache when live params or inputs changed."""
        self._advance_bscan_floor_to_cpp_window()
        signature = self._bscan_signature()
        if signature == self._bscan_render_signature:
            return
        self._rebuild_bscan_history_from_results()
        self._bscan_render_signature = signature

    def _bscan_signature(self) -> tuple[object, ...]:
        """Build state signature for B-scan history cache invalidation."""
        live_config = self._live_processing_config()
        result_history = list(self._result_history)
        return build_bscan_signature(
            live_config=live_config,
            subtract_mean_ascan_enabled=bool(self._bscan_subtract_mean_ascan.isChecked()),
            result_history=result_history,
            history_limit=self._bscan_history_limit,
            floor_collection_id=self._bscan_history_floor_collection_id,
        )

    def _rebuild_bscan_history_from_results(self) -> None:
        """Recompute B-scan history cache from results history buffer."""
        result_history = list(self._result_history)
        history_by_combo, depth_axis_by_combo = rebuild_bscan_history_from_results(
            result_history=result_history,
            history_limit=self._bscan_history_limit,
            floor_collection_id=self._bscan_history_floor_collection_id,
        )
        self._bscan_history_by_combo = history_by_combo
        self._bscan_depth_axis_by_combo = depth_axis_by_combo

    def _pick_bscan_display_key(self) -> tuple[int, int] | None:
        """Choose combo history key to render."""
        requested_key = self._requested_bscan_display_key()
        display_key = pick_bscan_display_key(
            self._bscan_history_by_combo,
            preferred_key=requested_key,
        )
        available_keys = sorted(self._bscan_history_by_combo.keys())
        if requested_key is not None and requested_key != display_key and available_keys:
            combo_signature = ",".join(f"{input_pos}:{output_pos}" for input_pos, output_pos in available_keys)
            details = "\n".join(
                f"- in{input_pos}/out{output_pos}"
                for input_pos, output_pos in available_keys
            )
            self._log_warning(
                f"B-scan display combo in{requested_key[0]}/out{requested_key[1]} is not available; "
                f"rendering in{display_key[0]}/out{display_key[1]} instead.",
                details=details,
                once_key=(
                    f"bscan_requested_combo_missing_{requested_key[0]}_{requested_key[1]}_"
                    f"{combo_signature}"
                ),
            )
            return display_key
        if display_key is not None and len(available_keys) > 1:
            if requested_key == display_key:
                return display_key
            combo_signature = ",".join(f"{input_pos}:{output_pos}" for input_pos, output_pos in available_keys)
            details = "\n".join(
                f"- in{input_pos}/out{output_pos}"
                for input_pos, output_pos in available_keys
            )
            self._log(
                f"B-scan auto-selected combo in{display_key[0]}/out{display_key[1]} because multiple combos are available.",
                once_key=f"bscan_auto_display_{combo_signature}",
            )
            self._log_warning(
                "B-scan has multiple combo histories but the UI currently renders only one at a time.",
                details=details,
                once_key=f"bscan_multi_combo_warning_{combo_signature}",
            )
        return display_key

    def _requested_bscan_display_key(self) -> tuple[int, int] | None:
        """Return the multi-device B-scan display combo requested by the GUI."""
        if not (
            hasattr(self, "_is_multi_device_model_selected")
            and self._is_multi_device_model_selected()
            and hasattr(self, "_single_combo_input")
            and hasattr(self, "_single_combo_output")
        ):
            return None

        input_text = self._single_combo_input.text().strip()
        output_text = self._single_combo_output.text().strip()
        if not input_text or not output_text:
            return None

        try:
            return int(input_text), int(output_text)
        except ValueError:
            return None

    def _bscan_lookup_table(self, axis_mode: str) -> np.ndarray:
        """Return lookup table for current B-scan axis mode."""
        return bscan_lookup_table(axis_mode)

    @staticmethod
    def _build_lut(stops: list[str], *, size: int = 256) -> np.ndarray:
        """Backward-compatible wrapper around LUT builder."""
        return build_lut(stops, size=size)

    @staticmethod
    def _bscan_levels(sweeps: np.ndarray, axis_mode: str) -> tuple[float, float]:
        """Return display levels for B-scan image."""
        return bscan_levels(sweeps, axis_mode)

    def _clear_bscan_plot_history(self) -> None:
        """Drop cached B-scan history and invalidate cache signature."""
        self._bscan_history_by_combo.clear()
        self._bscan_depth_axis_by_combo.clear()
        self._bscan_render_signature = None

    def _advance_bscan_floor_to_cpp_window(self) -> None:
        """Clamp B-scan source history to C++ available replay window."""
        if not self._result_history:
            return

        cpp_window_limit = min(
            int(self._defaults_config.rings.preprocessed.capacity),
            int(self._defaults_config.rings.results.capacity),
        )
        cpp_window_limit = max(1, cpp_window_limit)
        latest_collection_id = int(self._result_history[-1].collection_id)
        current_floor = int(self._bscan_history_floor_collection_id)

        # Collection ids restart from 1 on new C++ run; release floor only while
        # acquisition is running, so manual "remove last" behavior in stopped mode
        # remains deterministic.
        if latest_collection_id < current_floor and self._supervisor.is_running():
            self._bscan_history_floor_collection_id = 0
            current_floor = 0

        floor_candidate = max(0, latest_collection_id - cpp_window_limit)
        if floor_candidate > current_floor:
            self._bscan_history_floor_collection_id = floor_candidate

    def _ensure_phase_view_box(self) -> pg.ViewBox:
        """Create or return secondary right-axis ViewBox for phase curves."""
        plot_item = self._bscan_plot.getPlotItem()
        phase_view_box = self._phase_viewbox
        if phase_view_box is None:
            phase_view_box = pg.ViewBox()
            self._phase_viewbox = phase_view_box
            plot_item.scene().addItem(phase_view_box)
            plot_item.getAxis("right").linkToView(phase_view_box)
            phase_view_box.setXLink(plot_item.vb)
            plot_item.vb.sigResized.connect(self._update_phase_view_box_geometry)
        self._update_phase_view_box_geometry()
        return phase_view_box

    def _update_phase_view_box_geometry(self) -> None:
        """Keep right-axis ViewBox geometry in sync with main plot ViewBox."""
        phase_view_box = self._phase_viewbox
        if phase_view_box is None:
            return
        plot_item = self._bscan_plot.getPlotItem()
        phase_view_box.setGeometry(plot_item.vb.sceneBoundingRect())
        phase_view_box.linkedViewChanged(plot_item.vb, phase_view_box.XAxis)

    def _clear_phase_overlay(self) -> None:
        """Remove all phase curves from secondary ViewBox."""
        self._trace_phase_plot.clear()

    def _disable_phase_axis(self) -> None:
        """Hide right axis and clear phase overlay when phase is not rendered."""
        self._clear_phase_overlay()
