# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
import io
import pathlib
import warnings
from collections.abc import Callable, Generator
from functools import partial, wraps
from typing import Any, TypeVar

import h5py
import numpy as np
import sciline as sl
import scipp as sc

from .types import (
    DetectorIndex,
    DetectorName,
    FilePath,
    NMXDetectorMetadata,
    NMXExperimentMetadata,
    NMXReducedDataGroup,
)


def _fallback_compute_positions(dg: sc.DataGroup) -> sc.DataGroup:
    import warnings

    import scippnexus as snx

    warnings.warn(
        "Using fallback compute_positions due to empty log entries. "
        "This may lead to incorrect results. Please check the data carefully."
        "The fallback will replace empty logs with a scalar value of zero.",
        UserWarning,
        stacklevel=2,
    )

    empty_transformations = [
        transformation
        for transformation in dg['depends_on'].transformations.values()
        if 'time' in transformation.value.dims
        and transformation.sizes['time'] == 0  # empty log
    ]
    for transformation in empty_transformations:
        orig_value = transformation.value
        orig_value = sc.scalar(0, unit=orig_value.unit, dtype=orig_value.dtype)
        transformation.value = orig_value
    return snx.compute_positions(dg, store_transform='transform_matrix')


def _compute_positions(
    dg: sc.DataGroup, auto_fix_transformations: bool = False
) -> sc.DataGroup:
    """Compute positions of the data group from transformations.

    Wraps the `scippnexus.compute_positions` function
    and provides a fallback for cases where the transformations
    contain empty logs.

    Parameters
    ----------
    dg:
        Data group containing the transformations and data.
    auto_fix_transformations:
        If `True`, it will attempt to fix empty transformations.
        It will replace them with a scalar value of zero.
        It is because adding a time dimension will make it not possible
        to compute positions of children due to time-dependent transformations.

    Returns
    -------
    :
        Data group with computed positions.

    Warnings
    --------
    If `auto_fix_transformations` is `True`, it will warn about the fallback
    being used due to empty logs or scalar transformations.
    This is because the fallback may lead to incorrect results.

    """
    import scippnexus as snx

    try:
        return snx.compute_positions(dg, store_transform='transform_matrix')
    except ValueError as e:
        if auto_fix_transformations:
            return _fallback_compute_positions(dg)
        raise e


def _create_dataset_from_string(*, root_entry: h5py.Group, name: str, var: str) -> None:
    root_entry.create_dataset(name, dtype=h5py.string_dtype(), data=var)


def _create_dataset_from_var(
    *,
    root_entry: h5py.Group,
    var: sc.Variable,
    name: str,
    long_name: str | None = None,
    compression: str | None = None,
    compression_opts: int | tuple[int, int] | None = None,
    chunks: tuple[int, ...] | int | bool | None = None,
    dtype: Any = None,
) -> h5py.Dataset:
    compression_options = {}
    if compression is not None:
        compression_options["compression"] = compression
    if compression_opts is not None:
        compression_options["compression_opts"] = compression_opts

    dataset = root_entry.create_dataset(
        name,
        data=var.values if dtype is None else var.values.astype(dtype, copy=False),
        chunks=chunks,
        **compression_options,
    )
    if var.unit is not None:
        dataset.attrs["units"] = str(var.unit)
    if long_name is not None:
        dataset.attrs["long_name"] = long_name
    return dataset


@wraps(_create_dataset_from_var)
def _create_compressed_dataset(*args, **kwargs):
    """Create dataset with compression options.

    It will try to use ``bitshuffle`` for compression if available.
    Otherwise, it will fall back to ``gzip`` compression.

    [``Bitshuffle/LZ4``](https://github.com/kiyo-masui/bitshuffle)
    is used for convenience.
    Since ``Dectris`` uses it for their Nexus file compression,
    it is compatible with DIALS.
    ``Bitshuffle/LZ4`` tends to give similar results to
    GZIP and other compression algorithms with better performance.
    A naive implementation of bitshuffle/LZ4 compression,
    shown in [issue #124](https://github.com/scipp/essnmx/issues/124),
    led to 80% file reduction (365 MB vs 1.8 GB).

    """
    try:
        import bitshuffle.h5

        compression_filter = bitshuffle.h5.H5FILTER
        default_compression_opts = (0, bitshuffle.h5.H5_COMPRESS_LZ4)
    except ImportError:
        warnings.warn(
            UserWarning(
                "Could not find the bitshuffle.h5 module from bitshuffle package. "
                "The bitshuffle package is not installed or only partially installed. "
                "Exporting to NeXus files with bitshuffle compression is not possible."
            ),
            stacklevel=2,
        )
        compression_filter = "gzip"
        default_compression_opts = 4

    return _create_dataset_from_var(
        *args,
        **kwargs,
        compression=compression_filter,
        compression_opts=default_compression_opts,
    )


def _create_root_data_entry(file_obj: h5py.File) -> h5py.Group:
    nx_entry = file_obj.create_group("NMX_data")
    nx_entry.attrs["NX_class"] = "NXentry"
    nx_entry.attrs["default"] = "data"
    nx_entry.attrs["name"] = "NMX"
    nx_entry["name"] = "NMX"
    nx_entry["definition"] = "TOFRAW"
    return nx_entry


def _create_sample_group(data: sc.DataGroup, nx_entry: h5py.Group) -> h5py.Group:
    nx_sample = nx_entry.create_group("NXsample")
    nx_sample["name"] = data['sample_name'].value
    _create_dataset_from_var(
        root_entry=nx_sample,
        var=data['crystal_rotation'],
        name='crystal_rotation',
        long_name='crystal rotation in Phi (XYZ)',
    )
    return nx_sample


def _create_instrument_group(data: sc.DataGroup, nx_entry: h5py.Group) -> h5py.Group:
    nx_instrument = nx_entry.create_group("NXinstrument")
    nx_instrument.create_dataset("proton_charge", data=data['proton_charge'].values)

    nx_detector_1 = nx_instrument.create_group("detector_1")
    # Detector counts
    _create_compressed_dataset(
        root_entry=nx_detector_1,
        name="counts",
        var=data['counts'],
    )
    # Time of arrival bin edges
    _create_dataset_from_var(
        root_entry=nx_detector_1,
        var=data['counts'].coords['t'],
        name="t_bin",
        long_name="t_bin TOF (ms)",
    )
    # Pixel IDs
    _create_compressed_dataset(
        root_entry=nx_detector_1,
        name="pixel_id",
        var=data['counts'].coords['id'],
        long_name="pixel ID",
    )
    return nx_instrument


def _create_detector_group(data: sc.DataGroup, nx_entry: h5py.Group) -> h5py.Group:
    nx_detector = nx_entry.create_group("NXdetector")
    # Position of the first pixel (lowest ID) in the detector
    _create_compressed_dataset(
        root_entry=nx_detector,
        name="origin",
        var=data['origin_position'],
    )
    # Fast axis, along where the pixel ID increases by 1
    _create_dataset_from_var(
        root_entry=nx_detector, var=data['fast_axis'], name="fast_axis"
    )
    # Slow axis, along where the pixel ID increases
    # by the number of pixels in the fast axis
    _create_dataset_from_var(
        root_entry=nx_detector, var=data['slow_axis'], name="slow_axis"
    )
    return nx_detector


def _create_source_group(data: sc.DataGroup, nx_entry: h5py.Group) -> h5py.Group:
    nx_source = nx_entry.create_group("NXsource")
    nx_source["name"] = "European Spallation Source"
    nx_source["short_name"] = "ESS"
    nx_source["type"] = "Spallation Neutron Source"
    nx_source["distance"] = sc.norm(data['source_position']).value
    nx_source["probe"] = "neutron"
    nx_source["target_material"] = "W"
    return nx_source


def export_as_nexus(
    data: sc.DataGroup, output_file: str | pathlib.Path | io.BytesIO
) -> None:
    """Export the reduced data to a NeXus file.

    Currently exporting step is not expected to be part of sciline pipelines.
    """
    warnings.warn(
        DeprecationWarning(
            "Exporting to custom NeXus format will be deprecated in the near future "
            ">=26.12.0. "
            "Please use ``export_as_nxlauetof`` instead."
        ),
        stacklevel=2,
    )
    with h5py.File(output_file, "w") as f:
        f.attrs["default"] = "NMX_data"
        nx_entry = _create_root_data_entry(f)
        _create_sample_group(data, nx_entry)
        _create_instrument_group(data, nx_entry)
        _create_detector_group(data, nx_entry)
        _create_source_group(data, nx_entry)


def _create_lauetof_data_entry(file_obj: h5py.File) -> h5py.Group:
    nx_entry = file_obj.create_group("entry")
    nx_entry.attrs["NX_class"] = "NXentry"
    return nx_entry


def _add_lauetof_definition(nx_entry: h5py.Group) -> None:
    _create_dataset_from_string(root_entry=nx_entry, name="definition", var="NXlauetof")


def _add_lauetof_instrument(nx_entry: h5py.Group) -> h5py.Group:
    nx_instrument = nx_entry.create_group("instrument")
    nx_instrument.attrs["NX_class"] = "NXinstrument"
    _create_dataset_from_string(root_entry=nx_instrument, name="name", var="NMX")
    return nx_instrument


def _add_lauetof_source_group(
    dg: NMXExperimentMetadata, nx_instrument: h5py.Group
) -> None:
    nx_source = nx_instrument.create_group("source")
    nx_source.attrs["NX_class"] = "NXsource"
    _create_dataset_from_string(
        root_entry=nx_source, name="name", var="European Spallation Source"
    )
    _create_dataset_from_string(root_entry=nx_source, name="short_name", var="ESS")
    _create_dataset_from_string(
        root_entry=nx_source, name="type", var="Spallation Neutron Source"
    )
    _create_dataset_from_var(
        root_entry=nx_source, name="distance", var=sc.norm(dg["source_position"])
    )
    # Legacy probe information.
    _create_dataset_from_string(root_entry=nx_source, name="probe", var="neutron")


def _add_lauetof_detector_group(dg: sc.DataGroup, nx_instrument: h5py.Group) -> None:
    nx_detector = nx_instrument.create_group(dg["detector_name"].value)  # Detector name
    nx_detector.attrs["NX_class"] = "NXdetector"
    _create_dataset_from_var(
        name="polar_angle",
        root_entry=nx_detector,
        var=sc.scalar(0, unit='deg'),  # TODO: Add real data
    )
    _create_dataset_from_var(
        name="azimuthal_angle",
        root_entry=nx_detector,
        var=sc.scalar(0, unit='deg'),  # TODO: Add real data
    )
    _create_dataset_from_var(
        name="x_pixel_size", root_entry=nx_detector, var=dg["x_pixel_size"]
    )
    _create_dataset_from_var(
        name="y_pixel_size", root_entry=nx_detector, var=dg["y_pixel_size"]
    )
    _create_dataset_from_var(
        name="distance",
        root_entry=nx_detector,
        var=sc.scalar(0, unit='m'),  # TODO: Add real data
    )
    # Legacy geometry information until we have a better way to store it
    _create_dataset_from_var(
        name="origin", root_entry=nx_detector, var=dg['origin_position']
    )
    # Fast axis, along where the pixel ID increases by 1
    _create_dataset_from_var(
        root_entry=nx_detector, var=dg['fast_axis'], name="fast_axis"
    )
    # Slow axis, along where the pixel ID increases
    # by the number of pixels in the fast axis
    _create_dataset_from_var(
        root_entry=nx_detector, var=dg['slow_axis'], name="slow_axis"
    )


def _add_lauetof_sample_group(dg: NMXExperimentMetadata, nx_entry: h5py.Group) -> None:
    nx_sample = nx_entry.create_group("sample")
    nx_sample.attrs["NX_class"] = "NXsample"
    _create_dataset_from_var(
        root_entry=nx_sample,
        var=dg['crystal_rotation'],
        name='crystal_rotation',
        long_name='crystal rotation in Phi (XYZ)',
    )
    _create_dataset_from_string(
        root_entry=nx_sample,
        name='name',
        var=dg['sample_name'].value,
    )
    _create_dataset_from_var(
        name='orientation_matrix',
        root_entry=nx_sample,
        var=sc.array(
            dims=['i', 'j'],
            values=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
            unit="dimensionless",
        ),  # TODO: Add real data, the sample orientation matrix
    )
    _create_dataset_from_var(
        name='unit_cell',
        root_entry=nx_sample,
        var=sc.array(
            dims=['i'],
            values=[1.0, 1.0, 1.0, 90.0, 90.0, 90.0],
            unit="dimensionless",  # TODO: Add real data,
            # a, b, c, alpha, beta, gamma
        ),
    )


def _add_lauetof_monitor_group(data: sc.DataGroup, nx_entry: h5py.Group) -> None:
    nx_monitor = nx_entry.create_group("control")
    nx_monitor.attrs["NX_class"] = "NXmonitor"
    _create_dataset_from_string(root_entry=nx_monitor, name='mode', var='monitor')
    nx_monitor["preset"] = 0.0  # Check if this is the correct value
    data_dset = _create_dataset_from_var(
        name='data',
        root_entry=nx_monitor,
        var=sc.array(
            dims=['tof'], values=[1, 1, 1], unit="counts"
        ),  # TODO: Add real data, bin values
    )
    data_dset.attrs["signal"] = 1
    data_dset.attrs["primary"] = 1
    _create_dataset_from_var(
        name='time_of_flight',
        root_entry=nx_monitor,
        var=sc.array(
            dims=['tof'], values=[1, 1, 1], unit="s"
        ),  # TODO: Add real data, bin edges
    )


def _add_arbitrary_metadata(
    nx_entry: h5py.Group, **arbitrary_metadata: sc.Variable
) -> None:
    if not arbitrary_metadata:
        return

    metadata_group = nx_entry.create_group("metadata")
    for key, value in arbitrary_metadata.items():
        if not isinstance(value, sc.Variable):
            import warnings

            msg = f"Skipping metadata key '{key}' as it is not a scipp.Variable."
            warnings.warn(UserWarning(msg), stacklevel=2)
            continue
        else:
            _create_dataset_from_var(
                name=key,
                root_entry=metadata_group,
                var=value,
            )


def _export_static_metadata_as_nxlauetof(
    experiment_metadata: NMXExperimentMetadata,
    output_file: str | pathlib.Path | io.BytesIO,
    **arbitrary_metadata: sc.Variable,
) -> None:
    """Export the metadata to a NeXus file with the LAUE_TOF application definition.

    ``Metadata`` in this context refers to the information
    that is not part of the reduced detector counts itself,
    but is necessary for the interpretation of the reduced data.
    Since NMX can have arbitrary number of detectors,
    this function can take multiple detector metadata objects.

    Parameters
    ----------
    experiment_metadata:
        Experiment metadata object.
    output_file:
        Output file path.
    arbitrary_metadata:
        Arbitrary metadata that does not fit into the existing metadata objects.

    """
    with h5py.File(output_file, "w") as f:
        f.attrs["NX_class"] = "NXlauetof"
        nx_entry = _create_lauetof_data_entry(f)
        _add_lauetof_definition(nx_entry)
        _add_lauetof_sample_group(experiment_metadata, nx_entry)
        nx_instrument = _add_lauetof_instrument(nx_entry)
        _add_lauetof_source_group(experiment_metadata, nx_instrument)
        # Placeholder for ``monitor`` group
        _add_lauetof_monitor_group(experiment_metadata, nx_entry)
        # Skipping ``NXdata``(name) field with data link
        # Add arbitrary metadata
        _add_arbitrary_metadata(nx_entry, **arbitrary_metadata)


def _export_detector_metadata_as_nxlauetof(
    *detector_metadatas: NMXDetectorMetadata,
    output_file: str | pathlib.Path | io.BytesIO,
    append_mode: bool = True,
) -> None:
    """Export the detector specific metadata to a NeXus file.

    Since NMX can have arbitrary number of detectors,
    this function can take multiple detector metadata objects.

    Parameters
    ----------
    detector_metadatas:
        Detector metadata objects.
    output_file:
        Output file path.

    """

    if not append_mode:
        raise NotImplementedError("Only append mode is supported for now.")

    with h5py.File(output_file, "r+") as f:
        nx_entry = f["entry"]
        if "instrument" not in nx_entry:
            nx_instrument = _add_lauetof_instrument(f["entry"])
        else:
            nx_instrument = nx_entry["instrument"]
        # Add detector group metadata
        for detector_metadata in detector_metadatas:
            _add_lauetof_detector_group(detector_metadata, nx_instrument)


def _extract_counts(dg: sc.DataGroup) -> sc.Variable:
    counts: sc.DataArray = dg['counts'].data
    if 'id' in counts.dims:
        num_x, num_y = dg["detector_shape"].value
        return sc.fold(counts, dim='id', sizes={'x': num_x, 'y': num_y})
    else:
        # If there is no 'id' dimension, we assume it is already in the correct shape
        return counts


def _export_reduced_data_as_nxlauetof(
    dg: NMXReducedDataGroup,
    output_file: str | pathlib.Path | io.BytesIO,
    *,
    append_mode: bool = True,
    compress_counts: bool = True,
) -> None:
    """Export the reduced data to a NeXus file with the LAUE_TOF application definition.

    Even though this function only exports
    reduced data(detector counts and its coordinates),
    the input should contain all the necessary metadata
    for minimum sanity check.

    Parameters
    ----------
    dg:
        Reduced data and metadata.
    output_file:
        Output file path.
    append_mode:
        If ``True``, the file is opened in append mode.
        If ``False``, the file is opened in None-append mode.
        > None-append mode is not supported for now.
        > Only append mode is supported for now.
    compress_counts:
        If ``True``, the detector counts are compressed using bitshuffle.
        It is because only the detector counts are expected to be large.

    """
    if not append_mode:
        raise NotImplementedError("Only append mode is supported for now.")

    with h5py.File(output_file, "r+") as f:
        nx_detector: h5py.Group = f[f"entry/instrument/{dg['detector_name'].value}"]
        # Data - shape: [n_x_pixels, n_y_pixels, n_tof_bins]
        # The actual application definition defines it as integer,
        # but we keep the original data type for now
        num_x, num_y = dg["detector_shape"].value  # Probably better way to do this
        if compress_counts:
            data_dset = _create_compressed_dataset(
                name="data",
                root_entry=nx_detector,
                var=_extract_counts(dg),
                chunks=(num_x, num_y, 1),
                dtype=np.uint,
            )
        else:
            data_dset = _create_dataset_from_var(
                name="data",
                root_entry=nx_detector,
                var=_extract_counts(dg),
                dtype=np.uint,
            )
        data_dset.attrs["signal"] = 1
        _create_dataset_from_var(
            name='time_of_flight',
            root_entry=nx_detector,
            var=sc.midpoints(dg['counts'].coords['t'], dim='t'),
        )


def _check_file(
    filename: str | pathlib.Path | io.BytesIO, overwrite: bool
) -> pathlib.Path | io.BytesIO:
    if isinstance(filename, str | pathlib.Path):
        filename = pathlib.Path(filename)
        if filename.exists() and not overwrite:
            raise FileExistsError(
                f"File '{filename}' already exists. Use `overwrite=True` to overwrite."
            )
    return filename


T = TypeVar("T", bound=sc.DataArray)


class NXLauetofWriter:
    def __init__(
        self,
        *,
        output_filename: str | pathlib.Path | io.BytesIO,
        workflow: sl.Pipeline,
        chunk_generator: Callable[[FilePath, DetectorName], Generator[T, None, None]],
        chunk_insert_key: type[T],
        extra_meta: dict[str, sc.Variable] | None = None,
        compress_counts: bool = True,
        overwrite: bool = False,
    ) -> None:
        from ess.reduce.streaming import EternalAccumulator, StreamProcessor

        from .types import FilePath, NMXReducedCounts

        self.compress_counts = compress_counts
        self._chunk_generator = chunk_generator
        self._chunk_insert_key = chunk_insert_key
        self._workflow = workflow
        self._output_filename = _check_file(output_filename, overwrite)
        self._input_filename = workflow.compute(FilePath)
        self._final_stream_processor = partial(
            StreamProcessor,
            dynamic_keys=(chunk_insert_key,),
            target_keys=(NMXReducedDataGroup,),
            accumulators={NMXReducedCounts: EternalAccumulator},
        )
        self._detector_metas: dict[DetectorName, NMXDetectorMetadata] = {}
        self._detector_reduced: dict[DetectorName, NMXReducedDataGroup] = {}
        _export_static_metadata_as_nxlauetof(
            experiment_metadata=self._workflow.compute(NMXExperimentMetadata),
            output_file=self._output_filename,
            **(extra_meta or {}),
        )

    def add_panel(
        self, *, detector_id: DetectorIndex | DetectorName
    ) -> NMXReducedDataGroup:
        from .types import PixelIds

        temp_wf = self._workflow.copy()
        if isinstance(detector_id, int):
            temp_wf[DetectorIndex] = detector_id
        elif isinstance(detector_id, str):
            temp_wf[DetectorName] = detector_id
        else:
            raise TypeError(
                f"Expected detector_id to be an int or str, got {type(detector_id)}"
            )

        _export_detector_metadata_as_nxlauetof(
            temp_wf.compute(NMXDetectorMetadata),
            output_file=self._output_filename,
        )
        # First compute static information
        detector_name = temp_wf.compute(DetectorName)
        temp_wf[PixelIds] = temp_wf.compute(PixelIds)
        processor = self._final_stream_processor(temp_wf)
        # Then iterate over the chunks
        for da in self._chunk_generator(self._input_filename, detector_name):
            if any(da.sizes.values()) == 0:
                continue
            else:
                results = processor.add_chunk({self._chunk_insert_key: da})

        _export_reduced_data_as_nxlauetof(
            results[NMXReducedDataGroup],
            self._output_filename,
            compress_counts=self.compress_counts,
        )
        return results[NMXReducedDataGroup]
