import logging
import warnings
from typing import Dict
from typing import List
from typing import Optional
from typing import Union

from ewokscore import Task
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from pydantic import Field

from ...io import hdf5
from ...io import output_uri
from ...xrffit.pymca_config import PyMcaXrfConfiguration
from ..positioner_utils import get_primary_beam_energy_value
from .execute import fit_multi
from .execute import fit_single

logger = logging.getLogger(__name__)


class FitSingleScanSingleDetectorInputs(BaseInputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    detector_name: str = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=["fx_nano_det0"],
    )
    output_root_uri: str = Field(
        description="Target HDF5 file URI with optional data path.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Optional group underneath ``output_root_uri``."
    )
    config: Union[str, PyMcaXrfConfiguration] = Field(
        description="PyMca configuration",
        examples=["/config/configuration.cfg", {"peaks": {}, "fit": {}}],
    )
    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )
    process_uri_template: str = Field(default=None, description="DEPRECATED")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None


class FitSingleScanSingleDetectorOutputs(BaseOutputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    output_root_uri: str = Field(
        description="Original output root URI received as input.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Original output root group received as input."
    )
    xrf_results_uri: str = Field(
        description="Output NXprocess results URI.",
        examples=["/results/dataset.h5::/1.1/fit/mca0/results"],
    )
    detector_name: str = Field(description="Name of the fitted detector.")


class FitSingleScanSingleDetector(
    Task,
    input_model=FitSingleScanSingleDetectorInputs,
    output_model=FitSingleScanSingleDetectorOutputs,
):
    """XRF fit of one scan with one detector"""

    def run(self):
        xrf_results_uri = fit_single(**self._fit_parameters())
        self.outputs.bliss_scan_uri = self.inputs.bliss_scan_uri
        self.outputs.detector_name = self.inputs.detector_name
        self.outputs.output_root_uri = self.inputs.output_root_uri
        self.outputs.output_root_group = self.inputs.output_root_group
        self.outputs.xrf_results_uri = xrf_results_uri

    def _fit_parameters(self) -> dict:
        params = self.get_input_values()

        params["pymca_configuration"] = params.pop("config")
        _resolve_energy(params)
        _resolve_output_uri(params)

        return params


class FitSingleScanMultiDetectorInputs(BaseInputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    detector_names: List[str] = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=[["fx_nano_det0", "fx_nano_det1"]],
    )
    output_root_uri: str = Field(
        description="Target HDF5 file URI with optional data path.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Optional group underneath ``output_root_uri``."
    )
    configs: List[Union[str, PyMcaXrfConfiguration]] = Field(
        description="PyMca configuration",
        examples=[
            ["/config/det0.cfg", "/config/det1.cfg"],
            [{"peaks": {}, "fit": {}}, {"peaks": {}, "fit": {}}],
        ],
    )
    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )
    process_uri_template: str = Field(default=None, description="DEPRECATED")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None


class FitSingleScanMultiDetectorOutputs(BaseOutputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    output_root_uri: str = Field(
        description="Original output root URI received as input.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Original output root group received as input."
    )
    xrf_results_uris: List[str] = Field(
        description="Output NXprocess results URI.",
        examples=[
            "/results/dataset.h5::/1.1/fit/fx_nano_det0/results",
            "/results/dataset.h5::/1.1/fit/fx_nano_det1/results",
        ],
    )
    detector_names: List[str] = Field(description="Names of the fitted detectors.")


class FitSingleScanMultiDetector(
    Task,
    input_model=FitSingleScanMultiDetectorInputs,
    output_model=FitSingleScanMultiDetectorOutputs,
):
    """XRF fit of one scan with multiple detectors"""

    def run(self):
        xrf_results_uris = fit_multi(**self._fit_parameters())
        self.outputs.bliss_scan_uri = self.inputs.bliss_scan_uri
        self.outputs.detector_names = self.inputs.detector_names
        self.outputs.output_root_uri = self.inputs.output_root_uri
        self.outputs.output_root_group = self.inputs.output_root_group
        self.outputs.xrf_results_uris = xrf_results_uris

    def _fit_parameters(self) -> dict:
        params = self.get_input_values()

        params["pymca_configurations"] = params.pop("configs")
        params["bliss_scan_uris"] = [params.pop("bliss_scan_uri")]
        _resolve_energies(params)
        _resolve_output_uri(params)

        return params


class FitStackSingleDetectorInputs(BaseInputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )
    detector_name: str = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=["fx_nano_det0"],
    )
    output_root_uri: str = Field(
        description="Target HDF5 file URI with optional data path.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Optional group underneath ``output_root_uri``."
    )
    config: Union[str, PyMcaXrfConfiguration] = Field(
        description="PyMca configuration",
        examples=["/config/configuration.cfg", {"peaks": {}, "fit": {}}],
    )
    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )
    process_uri_template: str = Field(default=None, description="DEPRECATED")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None


class FitStackSingleDetectorOutputs(BaseOutputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )
    output_root_uri: str = Field(
        description="Original output root URI received as input.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Original output root group received as input."
    )
    xrf_results_uri: str = Field(
        description="Output NXprocess results URI.",
        examples=["/results/dataset.h5::/1.1/fit/fx_nano_det0/results"],
    )
    detector_name: str = Field(description="Name of the fitted detector.")


class FitStackSingleDetector(
    Task,
    input_model=FitStackSingleDetectorInputs,
    output_model=FitStackSingleDetectorOutputs,
):
    """XRF fit of a stack of identical scan with one detector"""

    def run(self):
        xrf_results_uris = fit_multi(**self._fit_parameters())
        self.outputs.bliss_scan_uris = self.inputs.bliss_scan_uris
        self.outputs.detector_name = self.inputs.detector_name
        self.outputs.output_root_uri = self.inputs.output_root_uri
        self.outputs.output_root_group = self.inputs.output_root_group
        self.outputs.xrf_results_uri = xrf_results_uris[0]

    def _fit_parameters(self) -> dict:
        params = self.get_input_values()

        params["detector_names"] = [params.pop("detector_name")]
        params["pymca_configurations"] = [params.pop("config")]
        _resolve_energies(params)
        _resolve_output_uri(params)

        return params


class FitStackMultiDetectorInputs(BaseInputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )
    detector_names: List[str] = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=[["fx_nano_det0", "fx_nano_det1"]],
    )
    output_root_uri: str = Field(
        description="Target HDF5 file URI with optional data path.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Optional group underneath ``output_root_uri``."
    )
    configs: List[Union[str, PyMcaXrfConfiguration]] = Field(
        description="PyMca configuration",
        examples=[
            ["/config/det0.cfg", "/config/det1.cfg"],
            [{"peaks": {}, "fit": {}}, {"peaks": {}, "fit": {}}],
        ],
    )
    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )
    process_uri_template: str = Field(default=None, description="DEPRECATED")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None


class FitStackMultiDetectorOutputs(BaseOutputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )
    output_root_uri: str = Field(
        description="Original output root URI received as input.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Original output root group received as input."
    )
    xrf_results_uris: List[str] = Field(
        description="Output NXprocess results URI.",
        examples=[
            "/results/dataset.h5::/1.1/fit/fx_nano_det0/results",
            "/results/dataset.h5::/1.1/fit/fx_nano_det1/results",
        ],
    )
    detector_names: List[str] = Field(description="Names of the fitted detectors.")


class FitStackMultiDetector(
    Task,
    input_model=FitStackMultiDetectorInputs,
    output_model=FitStackMultiDetectorOutputs,
):
    """XRF fit of a stack of identical scan with multiple detectors"""

    def run(self):
        xrf_results_uris = fit_multi(**self._fit_parameters())
        self.outputs.bliss_scan_uris = self.inputs.bliss_scan_uris
        self.outputs.detector_names = self.inputs.detector_names
        self.outputs.output_root_uri = self.inputs.output_root_uri
        self.outputs.output_root_group = self.inputs.output_root_group
        self.outputs.xrf_results_uris = xrf_results_uris

    def _fit_parameters(self) -> dict:
        params = self.get_input_values()

        params["pymca_configurations"] = params.pop("configs")
        _resolve_energies(params)
        _resolve_output_uri(params)

        return params


def _resolve_energy(params: dict) -> None:
    energy_name = params.pop("energy_name", None)
    energy_uri_template = params.pop("energy_uri_template", None)
    params["energy"] = get_primary_beam_energy_value(
        params["bliss_scan_uri"],
        energy_name=energy_name,
        energy_uri_template=energy_uri_template,
        search_on_units=False,
    )


def _resolve_energies(params: dict) -> None:
    energy_name = params.pop("energy_name", None)
    energy_uri_template = params.pop("energy_uri_template", None)
    params["energies"] = [
        get_primary_beam_energy_value(
            bliss_scan_uri,
            energy_name=energy_name,
            energy_uri_template=energy_uri_template,
            search_on_units=False,
        )
        for bliss_scan_uri in params["bliss_scan_uris"]
    ]


def _resolve_output_uri(params: dict) -> None:
    output_root_uri = params.get("output_root_uri")
    output_root_group = params.pop("output_root_group", None)
    process_uri_template = params.pop("process_uri_template")

    if "bliss_scan_uri" in params:
        bliss_scan_uri = params.get("bliss_scan_uri")
        _, scan_h5path = hdf5.split_h5uri(bliss_scan_uri)
    else:
        bliss_scan_uris = params.get("bliss_scan_uris")
        _, scan_h5path = hdf5.split_h5uri(bliss_scan_uris[0])

    output_root_uri = output_uri.compose_full_output_uri(
        output_root_uri,
        default_output_data_path=scan_h5path,
        extra_data_paths=(output_root_group, output_uri.DEFAULT_FIT_NAME),
    )

    if process_uri_template:

        warnings.warn(
            "'process_uri_template' is deprecated and will be removed in a future version. ",
            DeprecationWarning,
            stacklevel=2,
        )

        logger.warning(
            "Ignore process_uri_template=%r and save fit results in %r",
            process_uri_template,
            output_root_uri,
        )

    params["output_root_uri"] = output_root_uri
