import importlib.metadata
from contextlib import ExitStack
from datetime import datetime
from typing import Dict
from typing import Optional

from ewoksdata.data.hdf5.dataset_writer import IndexedDatasetWriter
from ewoksdata.data.nexus import create_nexus_group
from ewoksdata.data.nexus import select_default_plot
from silx.io.dictdump import dicttonx

from ..fit.types import UNITS
from ..fit.types import XRFBatchFitResult
from ..pymca_config import PyMcaXrfConfiguration
from ..pymca_config import pymca_configdict_from_model


class XrfResultWriter:
    _PROGRAM_VERSION = importlib.metadata.version("ewoksfluo")

    def __init__(self, output_parent_uri: str, default_group: Optional[str] = None):
        self._output_parent_uri = output_parent_uri
        self._default_group = default_group

        # Output
        self._already_existed = False
        self._xrf_results_uri = None
        self._output_root_uri = None
        self._dataset_writers = {}

        # Context parameters
        self._reset_context_parameters()

    def _reset_context_parameters(self):
        self._ctx_stack = None
        self._nxprocess = None
        self._results = None
        self._hdf5_groups: Dict[str, dict] = dict()

    @property
    def xrf_results_uri(self) -> Optional[str]:
        return self._xrf_results_uri

    @property
    def output_root_uri(self) -> Optional[str]:
        return self._output_root_uri

    @property
    def already_existed(self) -> bool:
        return self._already_existed

    def _create_group(self, name: str, data: dict) -> None:
        is_nxdata = data["@NX_class"] == "NXdata"
        if is_nxdata:
            if self._results is None:
                self._results = self._nxprocess.create_group("results")
                self._results.attrs["NX_class"] = "NXcollection"
            parent = self._results
        else:
            parent = self._nxprocess
        if name in parent:
            return

        dicttonx(data, parent.file, h5path=f"{parent.name}/{name}")

        if is_nxdata:
            group = parent[name]
            self._hdf5_groups[name] = {"group": group, "signals": list()}
            if name == self._default_group:
                select_default_plot(group)

    def _create_nxdata_group(self, name: str, data: dict) -> None:
        if self._results is None:
            self._results = self._nxprocess.create_group("results")
            self._results.attrs["NX_class"] = "NXcollection"
        parent = self._results
        if name in parent:
            return

        dicttonx(data, parent.file, h5path=f"{parent.name}/{name}")

        group = parent[name]
        self._hdf5_groups[name] = {"group": group, "signals": list()}
        if name == self._default_group:
            select_default_plot(group)

    def save_configuration(self, configuration: PyMcaXrfConfiguration) -> None:
        if "configuration" not in self._hdf5_groups:
            data = {
                "@NX_class": "NXnote",
                "date": datetime.now().astimezone().isoformat(),
                "type": "text/plain",
                "data": pymca_configdict_from_model(configuration).tostring(),
            }
            self._create_group("configuration", data)

    def save_fit_result(self, result: XRFBatchFitResult, start: int, stop: int) -> None:
        for group in ["parameters", "uncertainties", "mass_fractions"]:
            if not result.has_data(group):
                continue
            if group not in self._hdf5_groups:
                data = {"@NX_class": "NXdata"}
                self._create_group(group, data)

            if group == "mass_fractions":
                parameter_names = result.mass_fraction_names
            else:
                parameter_names = result.parameter_names

            for name, values in zip(parameter_names, getattr(result, group)):
                writer = self._get_dataset_writer(group, name)
                writer.add_points(values, list(range(start, stop)))

        if result.has_data("mass_fractions_per_layer"):
            for i, mass_fractions in enumerate(
                result.mass_fractions_per_layer, start=1
            ):
                group = f"mass_fractions_layer{i}"

                if group not in self._hdf5_groups:
                    data = {"@NX_class": "NXdata"}
                    self._create_group(group, data)

                for name, values in zip(result.mass_fraction_names, mass_fractions):
                    writer = self._get_dataset_writer(group, name)
                    writer.add_points(values, list(range(start, stop)))

        if result.has_diagnostics:
            group = "fit"

            if group not in self._hdf5_groups:
                data = {
                    "@NX_class": "NXdata",
                    "@axes": [".", "energy"],
                    "energy": result.raw_energies,
                }
                _add_units_attribute(data, "raw_energies", "energy")
                self._create_group(group, data)
                attrs = {"interpretation": "spectrum"}
            else:
                attrs = None

            writer = self._get_dataset_writer(group, "model", attrs=attrs)
            writer.add_points(result.model, list(range(start, stop)))

            writer = self._get_dataset_writer(group, "data", attrs=attrs)
            writer.add_points(result.spectra, list(range(start, stop)))

            writer = self._get_dataset_writer(group, "residuals", attrs=attrs)
            writer.add_points(result.residuals, list(range(start, stop)))

        if result.has_data("derivatives"):
            if "derivatives" not in self._hdf5_groups:
                data = {
                    "@NX_class": "NXdata",
                    "@axes": ["energy"],
                    "energy": result.raw_energies,
                }
                _add_units_attribute(data, "raw_energies", "energy")
                for name, values in zip(result.parameter_names, result.derivatives):
                    data[name] = values
                    data[f"{name}@interpretation"] = "spectrum"
                self._create_group("derivatives", data)

        if result.has_data("quantification_parameters"):
            if "quantification" not in self._hdf5_groups:
                parameters = result.quantification_parameters._asdict()
                data = {"@NX_class": "NXparameters", **parameters}
                for name in parameters:
                    _add_units_attribute(data, name, name)
                self._create_group("quantification", data)

    def __enter__(self) -> "XrfResultWriter":
        self._ctx_stack = ExitStack().__enter__()
        ctx = create_nexus_group(
            self._output_parent_uri, default_levels=("results", "fit"), retry_timeout=0
        )
        self._nxprocess, self._already_existed = self._ctx_stack.enter_context(ctx)
        self._nxprocess.attrs["NX_class"] = "NXprocess"
        if not self._already_existed:
            self._nxprocess["program"] = "ewoksfluo"
            self._nxprocess["version"] = self._PROGRAM_VERSION
        self._xrf_results_uri = (
            f"{self._nxprocess.file.filename}::{self._nxprocess.name}/results"
        )
        entry_name = [s for s in self._nxprocess.parent.name.split("/") if s][0]
        self._output_root_uri = f"{self._nxprocess.file.filename}::/{entry_name}"
        return self

    def __exit__(self, *args) -> None:
        try:
            self._finalize()
        finally:
            try:
                result = self._ctx_stack.__exit__(*args)
            finally:
                self._reset_context_parameters()
        return result

    def _get_dataset_writer(
        self,
        group: str,
        name: str,
        npoints: Optional[int] = None,
        attrs: Optional[dict] = None,
    ) -> IndexedDatasetWriter:
        key = (group, name)
        if key in self._dataset_writers:
            return self._dataset_writers[key]

        name = name.replace(" ", "_")
        nxdata = self._hdf5_groups[group]

        ctx = IndexedDatasetWriter(
            parent=nxdata["group"], name=name, npoints=npoints, attrs=attrs
        )
        nxdata["signals"].append(name)
        writer = self._ctx_stack.enter_context(ctx)
        self._dataset_writers[key] = writer
        return writer

    def _finalize(self):
        for writer in self._dataset_writers.values():
            writer.flush_buffer()
        for nxdata in self._hdf5_groups.values():
            group = nxdata["group"]
            if "signal" in group.attrs:
                continue
            signals = nxdata["signals"]
            if signals:
                group.attrs["signal"] = signals[0]
            if len(signals) > 1:
                group.attrs["auxiliary_signals"] = signals[1:]


def _add_units_attribute(data: dict, model_name: str, hdf5_name: str):
    if UNITS.get(model_name):
        data[f"{hdf5_name}@units"] = UNITS[model_name]
