from typing import Dict
from typing import List
from typing import Optional
from typing import Union

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

from ...xrffit.pymca_config import PyMcaXrfConfiguration
from ..models import FluoInputs
from ..positioner_utils import get_primary_beam_energy_value
from .execute import fit_multi
from .execute import fit_single


class FitSingleScanSingleDetectorInputs(FluoInputs):
    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="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    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="fit/{}")
    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="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    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):
        params = self.get_input_values()
        params["pymca_configuration"] = params.pop("config")
        _parse_energy(params)
        xrf_results_uri, output_root_uri = fit_single(**params)
        self.outputs.xrf_results_uri = xrf_results_uri
        self.outputs.bliss_scan_uri = params["bliss_scan_uri"]
        self.outputs.detector_name = params["detector_name"]
        self.outputs.output_root_uri = output_root_uri


class FitSingleScanMultiDetectorInputs(FluoInputs):
    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="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    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="fit/{}")
    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="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    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):
        params = self.get_input_values()
        params["pymca_configurations"] = params.pop("configs")
        bliss_scan_uri = params.pop("bliss_scan_uri")
        params["bliss_scan_uris"] = [bliss_scan_uri]
        _parse_energies(params)
        xrf_results_uris, output_root_uri = fit_multi(**params)
        self.outputs.xrf_results_uris = xrf_results_uris
        self.outputs.bliss_scan_uri = bliss_scan_uri
        self.outputs.detector_names = params["detector_names"]
        self.outputs.output_root_uri = output_root_uri


class FitStackSingleDetectorInputs(FluoInputs):
    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="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    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="fit/{}")
    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="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    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):
        params = self.get_input_values()
        detector_name = params.pop("detector_name")
        params["detector_names"] = [detector_name]
        params["pymca_configurations"] = [params.pop("config")]
        _parse_energies(params)
        xrf_results_uris, output_root_uri = fit_multi(**params)
        self.outputs.xrf_results_uri = xrf_results_uris[0]
        self.outputs.bliss_scan_uris = params["bliss_scan_uris"]
        self.outputs.detector_name = detector_name
        self.outputs.output_root_uri = output_root_uri


class FitStackMultiDetectorInputs(FluoInputs):
    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="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    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="fit/{}")
    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="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    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):
        params = self.get_input_values()
        params["pymca_configurations"] = params.pop("configs")
        _parse_energies(params)
        xrf_results_uris, output_root_uri = fit_multi(**params)
        self.outputs.xrf_results_uris = xrf_results_uris
        self.outputs.bliss_scan_uris = params["bliss_scan_uris"]
        self.outputs.detector_names = params["detector_names"]
        self.outputs.output_root_uri = output_root_uri


def _parse_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 _parse_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"]
    ]
