from __future__ import annotations

from pathlib import Path

from ImageD11.unitcell import unitcell as UnitCell
from silx.gui import qt
from silx.gui.colors import Colormap
from silx.io.utils import DataUrl

from ewoks3dxrd.io import save_geometry_and_lattice_file
from ewoks3dxrd.nexus.peaks import read_peaks_attributes
from ewoks3dxrd.tasks.filter_by_lattice import FilterByLattice
from ewoks3dxrd.tasks.filter_by_lattice import Inputs

from .common.custom_peak_plot import CustomPeakPlot
from .common.ewoks3dxrd_widget import Ewoks3DXRDWidget
from .common.execution_control_group_box import ExecutionControlGroupBoxWithProcessName
from .common.peak_filter_plot2d import PeakFilterPlot2D
from .lattice.lattice_filter_control_box import LatticeControlBox


class OWLatticeFilter(Ewoks3DXRDWidget, ewokstaskclass=FilterByLattice):
    name = "Lattice Filter"
    description = "Filter geometry transformed, segmented 3D Peaks."
    icon = "icons/filters-1.svg"
    _ewoks_inputs_to_hide_from_orange = (
        "reciprocal_dist_max",
        "reciprocal_dist_tol",
        "overwrite",
        "lattice_file",
        "process_group_name",
    )

    def __init__(self):
        super().__init__()

        settingsWidget = qt.QWidget(self)
        settingLayout = qt.QVBoxLayout(settingsWidget)

        self._inputParametersBox = LatticeControlBox(self)
        self._executeBox = ExecutionControlGroupBoxWithProcessName(
            buttonText="Filter Peaks"
        )
        self._executeBox.sigExecuteRequested.connect(self.execute_with_validation)
        self._executeBox.sigUrlChanged.connect(self._onDataURLEditFinished)
        self._executeBox.setProcessName(
            Inputs.model_fields["process_group_name"].default
        )
        settingLayout.addWidget(self._inputParametersBox)
        settingLayout.addWidget(self._executeBox)
        self.addControlWidget(settingsWidget)

        self._inputParametersBox.latticePanel.sigPlotLatticeRings.connect(
            self._plotRings
        )
        self._ds_ring_limits = 1.5
        self._plot = PeakFilterPlot2D(self)
        self._customPlot = CustomPeakPlot(self)
        self._customPlot.sigAxisChanged.connect(self.plotCustomInPeaksPlot)
        self._customPlot.sigAxisChanged.connect(self.plotCustomFilteredPeaksPlot)
        splitter = qt.QSplitter(self)
        splitter.setOrientation(qt.Qt.Horizontal)
        splitter.addWidget(self._plot)
        splitter.addWidget(self._customPlot)

        self.addMainWidget(splitter)
        self.registerInput(
            "geometry_updated_data_url",
            self._executeBox.getInputURL,
            self._executeBox.setInputURL,
        )
        self.registerInput(
            "lattice_file",
            self._prepareLatticeFile,
            self._inputParametersBox.latticePanel.setLatticeFile,
        )
        self.registerInput(
            "reciprocal_dist_tol",
            self._inputParametersBox.latticePanel.getReciprocalDistTol,
            self._inputParametersBox.latticePanel.setReciprocalDistTol,
        )
        self.registerInput(
            "reciprocal_dist_max",
            self._inputParametersBox.latticePanel.getReciprocalDistMax,
            self._inputParametersBox.latticePanel.setReciprocalDistMax,
        )
        self.registerInput(
            "process_group_name",
            self._executeBox.getProcessName,
            self._executeBox.setProcessName,
        )
        self.registerInput(
            "overwrite",
            self._executeBox.isOverwriteChecked,
            self._executeBox.setOverwrite,
        )

    def __post_init__(self):
        super().__post_init__()
        self._plotIncomingPeaks()

    def _plotRings(self, limit: float | None = None):
        if limit is None:
            limit = self._ds_ring_limits
        else:
            self._ds_ring_limits = limit

        parms = self._inputParametersBox.latticePanel.getLatticeParameters()
        try:
            unit_cell = UnitCell(
                lattice_parameters=parms["lattice_parameters"],
                symmetry=int(parms["lattice_space_group"]),
            )

            self._plot.addRings(unit_cell, limit)
            self._plot.setGraphXLabel("Reciprocal Distance (ds)")
            self._plot.setGraphYLabel("Azimuthal Angle (eta)")
        except Exception as e:
            self.showError(e, title="Lattice Ring Plot")

    def handleNewSignals(self):
        input_peaks_url = self.get_task_input_value("geometry_updated_data_url", None)
        if input_peaks_url is not None:
            self._executeBox.setInputURL(input_peaks_url)
            self._plotIncomingPeaks()

    def _plotIncomingPeaks(self):
        input_url = self.get_task_input_value("geometry_updated_data_url", None)
        if input_url is None:
            return
        data_url = DataUrl(input_url)
        nexus_file_path, data_group_path = (
            data_url.file_path(),
            data_url.data_path(),
        )
        incoming_3d_peaks = read_peaks_attributes(
            filename=nexus_file_path,
            process_name=data_group_path,
        )
        availableKeys = [key for key in incoming_3d_peaks.keys()]
        self._customPlot.populateAxesToolbar(availableKeys)
        self._plot.addUnfilteredScatter(
            x=incoming_3d_peaks["ds"],
            y=incoming_3d_peaks["eta"],
            value=len(incoming_3d_peaks["eta"]) * [1],
            colormap=Colormap(name="viridis", autoscaleMode="percentile_1_99"),
        )
        self._plot.setGraphXLabel("Reciprocal distance (ds)")
        self._plot.setGraphYLabel("Sum Intensity")
        self._plot.resetZoomOnce()
        self.plotCustomInPeaksPlot()

    def validateInputs(self):
        self._inputParametersBox.validate()
        self._executeBox.validate()

    def handleSuccessfulExecution(self):
        if self.get_task_output_value("lattice_filtered_data_url", None) is None:
            return
        self.setOutputDataUrl(self.get_task_output_value("lattice_filtered_data_url"))
        self._plotTaskGeneratedOutput()

    def _plotTaskGeneratedOutput(self):
        data_url = DataUrl(self.getOutputDataUrl())
        nexus_file_path, lattice_data_group_path = (
            data_url.file_path(),
            data_url.data_path(),
        )
        lattice_3d_peaks = read_peaks_attributes(
            filename=nexus_file_path,
            process_name=lattice_data_group_path,
        )
        self._plotRings(limit=max(lattice_3d_peaks["ds"]))
        self._plot.addFilteredScatter(
            x=lattice_3d_peaks["ds"],
            y=lattice_3d_peaks["eta"],
            value=lattice_3d_peaks["Number_of_pixels"],
            colormap=Colormap(
                name="viridis", normalization="log", autoscaleMode="percentile_1_99"
            ),
        )
        self._plot.setGraphXLabel("Reciprocal distance (ds)")
        self._plot.setGraphYLabel("Azimuthal angle (eta)")
        self._plot.resetZoomOnce()
        self.plotCustomFilteredPeaksPlot()

    def _prepareLatticeFile(self) -> str:
        file_path = Path(DataUrl(self._executeBox.getInputURL()).file_path())
        params = self._inputParametersBox.latticePanel.getLatticeParameters()
        if params["lattice_name"] is None:
            raise ValueError(f"[{__name__}] Lattice name must be provided")
        if len(params["lattice_parameters"]) != 6:
            raise ValueError(f"[{__name__}] Lattice parameters are mis-aligned")
        if params["lattice_space_group"] is None:
            raise ValueError(f"[{__name__}] 'lattice_space_group' name missing")
        lattice_name: str = params["lattice_name"]
        lattice_file = file_path.parent / f"{lattice_name}.par"
        save_geometry_and_lattice_file(
            file_path=lattice_file,
            geom_dict={},
            lattice_dict={
                "lattice_parameters": params["lattice_parameters"],
                "lattice_space_group": params["lattice_space_group"],
            },
        )
        return str(lattice_file)

    def _onDataURLEditFinished(self) -> None:
        self._handlePeaksInputURL(
            urlWidget=self._executeBox.inputUrlWidget,
            defaultInputKey="geometry_updated_data_url",
            errorTitle="Invalid Incoming Geometry Data group path",
        )

    def plotCustomInPeaksPlot(self):
        input_url = self.get_task_input_value("geometry_updated_data_url", None)
        if input_url is None:
            return
        self._customPlot.plotUnfilteredPeaks(input_url)

    def plotCustomFilteredPeaksPlot(self):
        lattice_filtered_data_url = self.getOutputDataUrl()
        if lattice_filtered_data_url is None:
            return
        self._customPlot.plotFilteredPeaks(lattice_filtered_data_url)

    def displayOutputDataUrl(self):
        self._plotTaskGeneratedOutput()
        self.plotCustomFilteredPeaksPlot()
