from __future__ import annotations

import logging
import os
import threading
from typing import Optional

from ewoksorange.gui.orange_utils.signals import Input
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 read_lattice_cell_data
from ewoks3dxrd.nexus.peaks import read_peaks_attributes
from ewoks3dxrd.nexus.utils import group_exists
from ewoks3dxrd.tasks.grid_index_grains_sub_process import GridIndexGrainsSubProcess

from .common.ewoks3dxrd_grainplotter import Ewoks3DXRDGrainPlotter
from .common.execution_control_group_box import ExecutionControlGroupBox
from .common.grain_map_scene import GrainMapScene
from .common.peak_filter_plot2d import PeakFilterPlot2D
from .common.sphere import Spheres
from .common.sphere import build_grain_spheres
from .common.sphere import build_grains_from_ascii_file
from .indexer.grid_index_control_box import GridIndexControlBox

_logger = logging.getLogger(__name__)


class OWGridIndexing(
    Ewoks3DXRDGrainPlotter,
    ewokstaskclass=GridIndexGrainsSubProcess,
):
    name = "Grid Indexing"
    description = "Multiprocessing positional search for the grains."
    icon = "icons/grid-search.svg"
    _ewoks_inputs_to_hide_from_orange = (
        "overwrite",
        "grid_index_parameters",
        "grid_abs_x_limit",
        "grid_abs_y_limit",
        "grid_abs_z_limit",
        "grid_step",
        "seed",
        "analyse_folder",
    )

    class Inputs:
        lattice_file = Input("lattice_file", str)

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

        settingsWidget = qt.QWidget(self)
        settingLayout = qt.QVBoxLayout(settingsWidget)
        self._inputParametersBox = GridIndexControlBox(self)
        self._inputParametersBox.sigGridSettingsChanged.connect(
            self._checkExecuteButton
        )
        self._inputParametersBox.latticeFile.sigMasterFileChanged.connect(
            self.setLatticeFile
        )

        self._executeBox = ExecutionControlGroupBox(buttonText="Run Grid Indexing")
        self._executeBox.sigExecuteRequested.connect(self.execute_with_validation)
        self._executeBox.sigUrlChanged.connect(self._onDataURLEditFinished)
        settingLayout.addWidget(self._inputParametersBox, stretch=1)
        settingLayout.addWidget(self._executeBox)
        self.addControlWidget(settingsWidget)

        self._grainSphere: Optional[Spheres] = None

        self._grainsPlot = GrainMapScene(self)
        self._grainsPlot.grainSizeChanged.connect(self._updateGrainSize)

        self._peaksPlot = PeakFilterPlot2D(self)
        self._plotTabWidget = qt.QTabWidget()
        self._plotTabWidget.addTab(self._peaksPlot, "Incoming Peaks")
        self._plotTabWidget.addTab(self._grainsPlot, "Generated Grains")

        self.addMainWidget(self._plotTabWidget)

        self._latticeFile = None
        self._ds_ring_limits = 1.5

        self.registerInput(
            "indexer_filtered_data_url", self._getInputUrl, self.setInputPeaksDataURL
        )
        self.registerInput(
            "grid_index_parameters",
            self.getGridIndexParameters,
            self.setGridIndexParameters,
        )
        self.registerInput(
            "grid_abs_x_limit",
            self._inputParametersBox.getGridXLimit,
            self._inputParametersBox.setGridXLimit,
        )
        self.registerInput(
            "grid_abs_y_limit",
            self._inputParametersBox.getGridYLimit,
            self._inputParametersBox.setGridYLimit,
        )
        self.registerInput(
            "grid_abs_z_limit",
            self._inputParametersBox.getGridZLimit,
            self._inputParametersBox.setGridZLimit,
        )
        self.registerInput(
            "grid_step",
            self._inputParametersBox.getGridStep,
            self._inputParametersBox.setGridStep,
        )
        self.registerInput(
            "overwrite",
            self._executeBox.isOverwriteChecked,
            self._executeBox.setOverwrite,
        )

        self.task_executor.finished.connect(self._cleanupLogging)

    def handleNewSignals(self):
        input_peaks_url = self.get_task_input_value("indexer_filtered_data_url")
        if isinstance(input_peaks_url, str):
            self._executeBox.setInputURL(input_peaks_url)
            self._plotIncomingPeaks()
            self._checkExecuteButton()

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

    def execute_ewoks_task(self, log_missing_inputs=False):
        super().execute_ewoks_task(log_missing_inputs)
        task_proc_init_event: threading = (
            self.task_executor.current_task.get_task_init_event()
        )
        task_proc_init_event.wait(timeout=5)
        if self.task_exception:
            _logger.info(f"Task failed during setup: {self.task_exception}", "ERROR")
            self.showError(self.task_exception, "ERROR")
            return

        if self.task_executor.isRunning():
            self._setupLogging()
            self.grainTimerStart()
        else:
            _logger.info("Failed to start subprocess, ERROR")

    def _plotIncomingPeaks(self):
        self._noPeaksExpected = None
        input_peaks_url = self.get_task_input_value("indexer_filtered_data_url", None)
        if input_peaks_url is None:
            return

        data_url = DataUrl(input_peaks_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,
        )
        self._peaksPlot.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._ds_ring_limits = max(incoming_3d_peaks["ds"])

        self._peaksPlot.setGraphXLabel("Reciprocal distance (ds)")
        self._peaksPlot.setGraphYLabel("Sum Intensity")
        self._peaksPlot.resetZoom()
        if self._latticeFile:
            self._plotLatticeCurve()

    def _plotLatticeCurve(self):
        if self._latticeFile is None:
            return

        unit_cell = read_lattice_cell_data(lattice_data_file_path=self._latticeFile)
        unit_cell = UnitCell(
            lattice_parameters=unit_cell.lattice_parameters,
            symmetry=int(unit_cell.space_group),
        )
        self._peaksPlot.addRings(unit_cell, self._ds_ring_limits)

    def handleSuccessfulExecution(self):
        self.grainTimerStop()
        output_url = self.get_task_output_value("grid_indexed_grain_data_url", None)
        if output_url is None:
            return
        self.setOutputDataUrl(output_url)
        self._plotTabWidget.setCurrentWidget(self._grainsPlot)
        self.displayOutputDataUrl()

    def displayOutputDataUrl(self):
        _logger.debug("GridIndexGrain: Display Output Data Url")
        url = self.getOutputDataUrl()
        self._grainsPlot.getSceneWidget().clearItems()
        self._grainsPlot.getSceneWidget().resetZoom()
        self._grainSphere = build_grain_spheres(url, self._grainsPlot.getGrainSize())
        self._grainsPlot.getSceneWidget().addItem(self._grainSphere)

    @Inputs.lattice_file
    def setLatticeFile(self, lattice_file: str):
        if not os.path.exists(lattice_file):
            return
        self._latticeFile = lattice_file
        self._plotLatticeCurve()

    def _updateGrainSize(self, size: float):
        if self._grainSphere:
            self._grainSphere.setRadiiNorm(radiiNorm=size)

    def _getInputUrl(self) -> str:
        input_url = DataUrl(self._executeBox.getInputURL())
        group_exists(
            filename=input_url.file_path(),
            data_group_path=input_url.data_path(),
        )
        return input_url.path()

    def getGridIndexParameters(self) -> dict:
        return self._inputParametersBox.getGridIndexParameters()

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

    @qt.Slot(str)
    def _plotGrainsUpdate(self, grain_file_path: str):
        self._grainsPlot.getSceneWidget().clearItems()
        self._grainSphere = build_grains_from_ascii_file(
            grain_file_path, self._grainsPlot.getGrainSize()
        )
        if self._grainSphere:
            self._grainsPlot.getSceneWidget().addItem(self._grainSphere)

    def _checkExecuteButton(
        self,
    ):
        try:
            self._inputParametersBox.getGridIndexParameters()
        except ValueError:
            self._executeBox.executeBtn.setDisabled(True)
            return
        self._executeBox.executeBtn.setEnabled(
            self._executeBox.inputUrlWidget.validDataURL()
        )

    def setInputPeaksDataURL(self, dataURL: str):
        self._executeBox.setInputURL(dataURL)
        self._plotIncomingPeaks()
        self._checkExecuteButton()

    def setGridIndexParameters(self, gridConfig: dict):
        self._inputParametersBox.setGridIndexParameters(gridConfig)
        self._checkExecuteButton()
