from __future__ import annotations

import logging
import os
import threading
from typing import Optional

from ewokscore import missing_data
from silx.gui import qt
from silx.io.url import DataUrl

from ewoks3dxrd.models import SampleConfig
from ewoks3dxrd.nexus.grains import read_grains
from ewoks3dxrd.nexus.utils import group_exists
from ewoks3dxrd.tasks.make_grain_map_sub_process import MakeGrainMapSubProcess

from .common.ewoks3dxrd_grainplotter import Ewoks3DXRDGrainPlotter
from .common.execution_control_group_box import ExecutionGroup
from .common.grain_map_scene import GrainMapScene
from .common.sphere import Spheres
from .common.sphere import build_grain_spheres
from .common.sphere import build_grains_from_ascii_file
from .grain_stats_plot import GrainStatsPlot
from .makegrains.make_map_control_box import MakeMapControlBox

_logger = logging.getLogger(__name__)


class OWMakeGrainMap(Ewoks3DXRDGrainPlotter, ewokstaskclass=MakeGrainMapSubProcess):
    name = "Refine Grain Mapping"
    description = "Second Stage Indexing (refine grain positions)."
    icon = "icons/maps.svg"
    _ewoks_inputs_to_hide_from_orange = (
        "overwrite",
        "intensity_frac",
        "hkl_tols",
        "minpks",
        "intensity_two_theta_range",
        "symmetry",
        "analyse_folder",
        "process_group_name",
    )

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

        settingsWidget = qt.QWidget(self)
        settingLayout = qt.QVBoxLayout(settingsWidget)
        self._inputParametersBox = MakeMapControlBox(self)
        self._inputParametersBox.sampleConfigGroup.sigMasterFileChanged.connect(
            self._fillSampleFolderConfig
        )
        self._inputParametersBox.latticeFile.sigMasterFileChanged.connect(
            self._fillLatticeFilePath
        )
        self._executeBox = ExecutionGroup(buttonText="Run Grid Indexing")
        self._inputParametersBox.sigGrainUrlChanged.connect(self._fillGrainDataURL)
        self._inputParametersBox.sigStrongPeaksUrlChanged.connect(
            self._fillStrongFilterDataURL
        )
        self._inputParametersBox.sigFinePeaksUrlChanged.connect(
            self._fillAllFilterDataURL
        )
        self._executeBox.executeBtn.clicked.connect(self.execute_with_validation)
        settingLayout.addWidget(self._inputParametersBox, stretch=1)
        settingLayout.addWidget(self._executeBox)
        self.addControlWidget(settingsWidget)

        self._grainSphere: Optional[Spheres] = None
        self._refineGrainsPlot = GrainMapScene(self)
        self._plotTabWidget = qt.QTabWidget()
        self._plotTabWidget.addTab(self._refineGrainsPlot, "Refined Grains")
        self._grainStatsPlot = GrainStatsPlot(self)
        self._plotTabWidget.addTab(self._grainStatsPlot, "Grains Stats")
        self.addMainWidget(self._plotTabWidget)
        self._refineGrainsPlot.grainSizeChanged.connect(self._updateGrainSize)

        self.registerInput(
            "indexed_grain_data_url",
            self._inputParametersBox.getGrainURL,
            self._inputParametersBox.setGrainURL,
        )
        self.registerInput(
            "intensity_filtered_data_url",
            self._inputParametersBox.getStrongPeaksURL,
            self._inputParametersBox.setStrongPeaksURL,
        )
        self.registerInput(
            "hkl_tols",
            self._inputParametersBox.grainMapParameterGroup.getTolerances,
            self._inputParametersBox.grainMapParameterGroup.setTolerances,
        )
        self.registerInput(
            "minpks",
            self._inputParametersBox.grainMapParameterGroup.getMinPeaks,
            self._inputParametersBox.grainMapParameterGroup.setMinPeaks,
        )
        self.registerInput(
            "intensity_fine_filtered_data_url",
            self._getIntensityFineFilteredDataUrl,
            self._setIntensityFineFilteredDataUrl,
        )
        self.registerInput(
            "intensity_two_theta_range",
            self._inputParametersBox.grainMapParameterGroup.getTwoThetaRange,
            self._inputParametersBox.grainMapParameterGroup.setTwoThetaRange,
        )
        self.registerInput(
            "symmetry",
            self._inputParametersBox.grainMapParameterGroup.getSymmetry,
            self._inputParametersBox.grainMapParameterGroup.setSymmetry,
        )
        self.registerInput(
            "lattice_file",
            self._inputParametersBox.latticeFile.getText,
            self._inputParametersBox.latticeFile.setText,
        )
        self.registerInput(
            "overwrite",
            self._executeBox.isOverwriteChecked,
            self._executeBox.setOverwrite,
        )
        self.registerInput(
            "folder_file_config",
            self.getSampleConfig,
            self.setSampleConfig,
        )

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

    def _fillStrongFilterDataURL(
        self,
    ):
        self.set_default_input(
            "intensity_filtered_data_url", self._inputParametersBox.getStrongPeaksURL()
        )

    def _fillAllFilterDataURL(
        self,
    ):
        self.set_default_input(
            "intensity_fine_filtered_data_url",
            self._inputParametersBox.getFinePeaksURL(),
        )

    def _fillGrainDataURL(
        self,
    ):
        self.set_default_input(
            "indexed_grain_data_url", self._inputParametersBox.getGrainURL()
        )

    def _fillSampleFolderConfig(self, _: str):
        sampleConfig = (
            self._inputParametersBox.sampleConfigGroup.getSampleConfig().model_dump()
        )
        self.set_default_input("folder_file_config", sampleConfig)

    def _fillLatticeFilePath(self):
        latticeFilePath = self._inputParametersBox.latticeFile.getText()
        if not os.path.exists(latticeFilePath):
            return
        self.set_default_input("lattice_file", latticeFilePath)

    def _getIntensityFineFilteredDataUrl(self) -> missing_data.MissingData | str:
        if not self._inputParametersBox.finePeaksUrl.isChecked():
            return missing_data.MISSING_DATA

        url = self._inputParametersBox.getFinePeaksURL()
        if not url:
            return missing_data.MISSING_DATA

        return url

    def _setIntensityFineFilteredDataUrl(self, url: missing_data.MissingData | str):
        if not isinstance(url, str) or url == "":
            self._inputParametersBox.finePeaksUrl.setChecked(False)
            return

        self._inputParametersBox.finePeaksUrl.setChecked(True)
        self._inputParametersBox.setFinePeaksURL(url)

    def handleNewSignals(self):
        all_peaks_url = self.get_task_input_value(
            "intensity_fine_filtered_data_url", None
        )
        if isinstance(all_peaks_url, str) and all_peaks_url != "":
            self._inputParametersBox.finePeaksUrl.setChecked(True)
            self._inputParametersBox.setFinePeaksURL(all_peaks_url)

        strong_filtered_peaks_url = self.get_task_input_value(
            "intensity_filtered_data_url", None
        )
        if isinstance(strong_filtered_peaks_url, str):
            self._inputParametersBox.setStrongPeaksURL(strong_filtered_peaks_url)

        grains_data_url = self.get_task_input_value("indexed_grain_data_url", None)
        if isinstance(grains_data_url, str):
            self._inputParametersBox.setGrainURL(grains_data_url)

        lattice_file = self.get_task_input_value("lattice_file", None)
        if isinstance(lattice_file, str):
            self._inputParametersBox.latticeFile.setText(lattice_file)

        folder_file_config = self.get_task_input_value("folder_file_config", None)
        if isinstance(folder_file_config, dict):
            self.set_default_input("folder_file_config", folder_file_config)
            self._inputParametersBox.sampleConfigGroup.setSampleConfig(
                config=SampleConfig(**folder_file_config)
            )

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

    def execute_ewoks_task(self, log_missing_inputs=False):
        super().execute_ewoks_task(log_missing_inputs)
        task_proc_init_event: threading.Event = (
            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 handleSuccessfulExecution(self):
        self.grainTimerStop()
        if self.get_task_output_value("make_map_data_url", None) is None:
            return
        self.setOutputDataUrl(self.get_task_output_value("make_map_data_url"))
        self.displayOutputDataUrl()

    def displayOutputDataUrl(self):
        _logger.debug("MakeMapGrain: Display Output Data Url")
        url = self.getOutputDataUrl()
        self._refineGrainsPlot.getSceneWidget().clearItems()
        self._refineGrainsPlot.getSceneWidget().resetZoom()
        self._grainSphere = build_grain_spheres(
            url, self._refineGrainsPlot.getGrainSize()
        )
        self._refineGrainsPlot.getSceneWidget().addItem(self._grainSphere)
        data_url = DataUrl(url)
        group_exists(
            filename=data_url.file_path(), data_group_path=data_url.data_path()
        )
        grains = read_grains(
            grain_file_h5=data_url.file_path(),
            entry_name=data_url.data_path(),
            process_group_name="",
        )

        self._grainStatsPlot.plotStats(grains)

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

    def getSampleConfig(self) -> dict:
        return self._inputParametersBox.sampleConfigGroup.getSampleConfig().model_dump()

    def setSampleConfig(self, sampleConfig: dict):
        self._inputParametersBox.sampleConfigGroup.setSampleConfig(
            SampleConfig(**sampleConfig)
        )

    @qt.Slot(str)
    def _plotGrainsUpdate(self, grain_file_path: str):

        self._refineGrainsPlot.getSceneWidget().clearItems()
        self._grainSphere = build_grains_from_ascii_file(
            grain_file_path, self._refineGrainsPlot.getGrainSize()
        )
        if self._grainSphere:
            self._refineGrainsPlot.getSceneWidget().addItem(self._grainSphere)
