"""Flint-side widget for BM32 XRR analysis plots."""

from __future__ import annotations

import logging
from typing import Any
from typing import Dict
from typing import Tuple

import numpy
from silx.gui import qt
from silx.gui.plot import Plot1D

from ..flint import capture_errors

logger = logging.getLogger(__name__)


class Bm32XrrAnalysisWidget(qt.QWidget):
    def __init__(self, parent=None) -> None:
        super().__init__(parent=parent)
        self._plots: Dict[Tuple[int, int], Plot1D] = {}
        layout = qt.QGridLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(4)
        self._header = qt.QLabel(parent=self)
        self._header.setText("BM32 XRR")
        self._header.setAlignment(qt.Qt.AlignLeft | qt.Qt.AlignVCenter)
        self._header.setStyleSheet("font-weight: 400; font-size: 18px; padding: 4px;")
        layout.addWidget(self._header, 0, 0, 1, 3)
        for row in range(3):
            for col in range(3):
                plot = Plot1D(parent=self)
                plot.setGraphGrid(True)
                plot.setMinimumHeight(200)
                layout.addWidget(plot, row + 1, col)
                self._plots[(row, col)] = plot

    @capture_errors
    def clear_data(self) -> None:
        self._header.setText("BM32 XRR")
        for plot in self._plots.values():
            plot.clear()

    @capture_errors
    def set_result_data(
        self,
        result: dict,
        sample_name: str = "",
    ) -> None:
        if not isinstance(result, dict):
            logger.warning("No plottable XRR data found in workflow result")
            self.clear_data()
            self.setWindowTitle(sample_name or "BM32 XRR analysis")
            for plot in self._plots.values():
                plot.setGraphTitle("No plottable data in workflow result")
            return
        title = sample_name or "BM32 XRR analysis"
        self._header.setText(title)
        self.setWindowTitle(title)
        _draw_panels_from_result(self._plots, result)


def _as_array(value: Any) -> numpy.ndarray:
    arr = numpy.asarray(value)
    if arr.ndim == 0:
        return numpy.asarray([arr.item()])
    return arr.ravel()


def _get_series(result: Dict[str, Any], key: str) -> numpy.ndarray | None:
    value = result.get(key)
    if value is None:
        return None
    arr = _as_array(value)
    if arr.size <= 1:
        return None
    return arr


def _add_curve(
    plot: Plot1D, x: numpy.ndarray | None, y: numpy.ndarray | None, legend: str
):
    if x is None or y is None:
        return
    n = min(x.size, y.size)
    if n <= 0:
        return
    plot.addCurve(x[:n], y[:n], legend=legend)


def _draw_panels_from_result(
    plots: Dict[Tuple[int, int], Plot1D], result: Dict[str, Any]
):
    for plot in plots.values():
        plot.clear()

    tth = _get_series(result, "tth")
    tth_bg = _get_series(result, "tth_bg")
    det = _get_series(result, "det")
    det_bg = _get_series(result, "det_bg")
    detm = _get_series(result, "detm")
    detm_bg = _get_series(result, "detm_bg")
    tthc = _get_series(result, "tthc")
    tthc_bg = _get_series(result, "tthc_bg")
    tthcr = _get_series(result, "tthcr")
    tthcr_bg = _get_series(result, "tthcr_bg")
    detmn = _get_series(result, "detmn")
    detmn_bg = _get_series(result, "detmn_bg")
    detmnp = _get_series(result, "detmnp")
    detmnp_bg = _get_series(result, "detmnp_bg")
    detmnpr = _get_series(result, "detmnpr")
    detmnpr_bg = _get_series(result, "detmnpr_bg")
    detmnpri_bg = _get_series(result, "detmnpri_bg")
    fpc = _get_series(result, "fpc")
    q = _get_series(result, "q")
    r = _get_series(result, "R")
    qraw = _get_series(result, "qraw")
    rraw = _get_series(result, "Rraw")
    z = _get_series(result, "z")
    delta_rho = _get_series(result, "delta_rho")
    delta_rho_bl = _get_series(result, "delta_rho_bl")
    delta_rho_nc = _get_series(result, "delta_rho_nc")
    baseline = _get_series(result, "baseline")
    rq4_bl = _get_series(result, "Rq4_bl")
    rq4_sm = _get_series(result, "Rq4_sm")
    phase = _get_series(result, "phase")
    phase_bl = _get_series(result, "phase_bl")
    minima = _get_series(result, "minima")

    qc = None
    rc = None
    if q is not None and r is not None:
        mask = r > 0
        if numpy.any(mask):
            qc = q[mask]
            rc = r[mask]

    p = plots[(0, 0)]
    _add_curve(p, tth, det, "XRR")
    _add_curve(p, tth_bg if tth_bg is not None else tth, det_bg, "BG")
    p.setGraphTitle("Raw Detector Signal")
    p.setGraphXLabel("2θ (°)")
    p.setGraphYLabel("Detector (counts)")
    p.setYAxisLogarithmic(True)

    p = plots[(1, 0)]
    _add_curve(p, tth, detm, "XRR")
    _add_curve(p, tth_bg if tth_bg is not None else tth, detm_bg, "BG")
    p.setGraphTitle("Monitor-Normalized Signal")
    p.setGraphXLabel("2θ (°)")
    p.setGraphYLabel("I / I₀")
    p.setYAxisLogarithmic(True)

    p = plots[(2, 0)]
    _add_curve(p, tthc, detmn, "XRR")
    _add_curve(p, tthc_bg if tthc_bg is not None else tthc, detmn_bg, "BG")
    p.setGraphTitle("Direct-Beam Normalized")
    p.setGraphXLabel("2θ (°)")
    p.setGraphYLabel("I / (I₀·I_db)")
    p.setYAxisLogarithmic(True)

    p = plots[(0, 1)]
    _add_curve(p, tthc, detmn, "XRR-no fc")
    _add_curve(p, tthc, detmnp, "XRR")
    _add_curve(p, tthc_bg if tthc_bg is not None else tthc, detmnp_bg, "BG")
    _add_curve(p, tthcr, detmnpr, "XRR o-db")
    _add_curve(p, tthcr_bg if tthcr_bg is not None else tthcr, detmnpr_bg, "BG o-db")
    _add_curve(p, tthcr, detmnpri_bg, "BG interp")
    _add_curve(p, tthc, fpc, "fpc")
    p.setGraphTitle("Footprint Correction and BG Interpolation")
    p.setGraphXLabel("2θ (°)")
    p.setGraphYLabel("I / (I₀·I_db·f_p)")
    p.setYAxisLogarithmic(True)

    p = plots[(1, 1)]
    _add_curve(p, q, r, "R")
    _add_curve(p, qc, rc, "R>0")
    p.setGraphTitle("Reflectivity R(q)")
    p.setGraphXLabel("q (Å⁻¹)")
    p.setGraphYLabel("R")
    p.setYAxisLogarithmic(True)

    p = plots[(2, 1)]
    if qraw is not None and rraw is not None:
        _add_curve(p, qraw, rraw * qraw**4, "raw")
    if qc is not None and rc is not None:
        _add_curve(p, qc, rc * qc**4, "R>0")
    p.setGraphTitle("q⁴R (Raw)")
    p.setGraphXLabel("q (Å⁻¹)")
    p.setGraphYLabel("q⁴R")

    p = plots[(0, 2)]
    if qc is not None and rc is not None:
        _add_curve(p, qc, rc * qc**4, "q^4R>0")
    _add_curve(p, q, rq4_bl, "bl")
    _add_curve(p, q, rq4_sm, "sm_bl")
    _add_curve(p, q, baseline, "baseline")
    if q is not None and rq4_sm is not None and minima is not None:
        idx = minima.astype(int)
        idx = idx[(idx >= 0) & (idx < min(q.size, rq4_sm.size))]
        if idx.size:
            p.addCurve(
                q[idx],
                rq4_sm[idx],
                legend="minima",
                symbol="o",
                linestyle="",
            )
    p.setGraphTitle("Baseline and Minima")
    p.setGraphXLabel("q (Å⁻¹)")
    p.setGraphYLabel("q⁴R")

    p = plots[(1, 2)]
    if q is not None and r is not None:
        _add_curve(p, q, numpy.sqrt(numpy.clip(r * q**4, 0, None)), "nc")
    _add_curve(p, q, phase, "phase")
    _add_curve(p, q, phase_bl, "phase+bl")
    p.setGraphTitle("Phased Signal √(q⁴R)")
    p.setGraphXLabel("q (Å⁻¹)")
    p.setGraphYLabel("√(q⁴R)")

    p = plots[(2, 2)]
    _add_curve(p, z, delta_rho_nc, "nc")
    _add_curve(p, z, delta_rho, "inverted")
    _add_curve(p, z, delta_rho_bl, "inverted bl")
    p.setGraphTitle("Electron Density Profile")
    p.setGraphXLabel("z (Å)")
    p.setGraphYLabel("Δρ (e⁻·Å⁻³)")
