from typing import List

import numpy
from ewokscore import Task
from ewokscore.model import BaseOutputModel
from ewoksdata.data.hdf5.config import guess_dataset_config
from pydantic import Field

from ...io.hdf5 import split_h5uri
from .. import nexus_utils
from ..hdf5_utils import create_hdf5_link
from ..hdf5_utils import link_bliss_scan
from ..models import FluoInputs
from .sum_utils import sum_spectra_from_hdf5


class Inputs(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/sumspectra",
        ],
    )

    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )

    detector_normalization_template: str = Field(
        default="1./<instrument/{}/live_time>",
        description="Expression to evaluate and multiply each spectrum with before summing. "
        "Templates marked with '<...>' are relative to `bliss_scan_uri`.",
    )

    output_detector_name: str = Field(default="mcasum")


class Outputs(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"]
    )
    detector_name: str = Field(
        description="Name of the detector containing the summed XRF spectra."
    )
    xrf_spectra_uri_template: str = Field(
        description="HDF5 group URI template of an XRF detector relative to `output_root_uri`.",
    )


class SumXrfSpectra(Task, input_model=Inputs, output_model=Outputs):
    """Add single-scan XRF spectra from multiple detectors."""

    def run(self):
        start_time = nexus_utils.now()

        _, scan_h5path = split_h5uri(self.inputs.bliss_scan_uri)

        sumdetector_name = self.inputs.output_detector_name

        with nexus_utils.save_in_ewoks_process(
            self.inputs.output_root_uri,
            start_time,
            process_config={
                "detector_normalization_template": self.inputs.detector_normalization_template
            },
            default_levels=(scan_h5path, "sumspectra"),
        ) as (process_group, already_existed):
            outentry = process_group.parent
            if not already_existed:
                sum_spectra: numpy.ndarray = sum_spectra_from_hdf5(
                    self.inputs.bliss_scan_uri,
                    self.inputs.xrf_spectra_uri_template,
                    self.inputs.detector_normalization_template,
                    self.inputs.detector_names,
                )

                link_bliss_scan(outentry, self.inputs.bliss_scan_uri, retry_timeout=0)

                nxdata = nexus_utils.create_nxdata(
                    process_group, "mcasum", signal="data"
                )

                dataset_kwargs = guess_dataset_config(
                    scan_shape=(sum_spectra.shape[0],),
                    detector_shape=(sum_spectra.shape[1],),
                    dtype=sum_spectra.dtype,
                )
                dset = nxdata.create_dataset("data", data=sum_spectra, **dataset_kwargs)
                dset.attrs["interpretation"] = "spectrum"

                nxdetector = outentry["instrument"].create_group(sumdetector_name)
                nxdetector.attrs["NX_class"] = "NXdetector"
                create_hdf5_link(nxdetector, "data", dset)
                create_hdf5_link(outentry["measurement"], sumdetector_name, dset)

            output_root_uri = f"{outentry.file.filename}::{outentry.name}"

        self.outputs.bliss_scan_uri = output_root_uri
        self.outputs.detector_name = sumdetector_name
        self.outputs.xrf_spectra_uri_template = self.inputs.xrf_spectra_uri_template
        self.outputs.output_root_uri = output_root_uri
