# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
# @author Jan-Lukas Wynen

from typing import Callable, Dict, Tuple, Union

import scipp as sc

from ..conversion import graph as _graphs


def _inelastic_scatter_graph(energy_mode):
    inelastic_graph_factory = {
        'direct_inelastic': _graphs.tof.direct_inelastic,
        'indirect_inelastic': _graphs.tof.indirect_inelastic,
    }
    return {
        **_graphs.beamline.beamline(scatter=True),
        **inelastic_graph_factory[energy_mode](start='tof'),
    }


def _reachable_by(target, graph):
    return any(
        target == targets if isinstance(targets, str) else target in targets
        for targets in graph.keys()
    )


def _elastic_scatter_graph(origin, target):
    scatter_graph_kinematics = _graphs.beamline.beamline(scatter=True)
    if _reachable_by(target, scatter_graph_kinematics):
        return dict(scatter_graph_kinematics)
    return {**scatter_graph_kinematics, **_graphs.tof.elastic(origin)}


def _scatter_graph(origin, target, energy_mode):
    graph = (
        _elastic_scatter_graph(origin, target)
        if energy_mode == 'elastic'
        else _inelastic_scatter_graph(energy_mode)
    )
    return graph


def conversion_graph(
    origin: str, target: str, scatter: bool, energy_mode: str
) -> Dict[Union[str, Tuple[str]], Callable]:
    """
    Get a conversion graph for given parameters.

    The graph can be used with `scipp.transform_coords`.

    :param origin: Name of the input coordinate.
    :param target: Name of the output coordinate.
    :param scatter: Choose whether to use scattering or non-scattering conversions.
    :param energy_mode: Select if energy is conserved. One of `'elastic'`,
                        `'direct_inelastic'`, `'indirect_inelastic'`.
    :return: Conversion graph.
    :seealso: :py:func:`scippneutron.convert`,
              :py:func:`scippneutron.deduce_conversion_graph`.
    """

    # Results are copied to ensure users do not modify the global dictionaries.
    if scatter:
        return dict(_scatter_graph(origin, target, energy_mode))
    else:
        return {
            **_graphs.beamline.beamline(scatter=False),
            **_graphs.tof.kinematic(start='tof'),
        }


def _find_inelastic_inputs(data):
    return [name for name in ('incident_energy', 'final_energy') if name in data.coords]


def _deduce_energy_mode(data, origin, target):
    inelastic_inputs = _find_inelastic_inputs(data)
    if target == 'energy_transfer':
        if len(inelastic_inputs) > 1:
            raise RuntimeError(
                "Data contains coords for incident *and* final energy, cannot have "
                "both for inelastic scattering."
            )
        if len(inelastic_inputs) == 0:
            raise RuntimeError(
                "Data contains neither coords for incident nor for final energy, this "
                "does not appear to be inelastic-scattering data, cannot convert to "
                "energy transfer."
            )
        return {
            'incident_energy': 'direct_inelastic',
            'final_energy': 'indirect_inelastic',
        }[inelastic_inputs[0]]

    if 'energy' in (origin, target):
        if inelastic_inputs:
            raise RuntimeError(
                f"Data contains coords for inelastic scattering "
                f"({inelastic_inputs}) but conversion with elastic energy requested. "
                f"This is not implemented."
            )
    return 'elastic'


def deduce_conversion_graph(
    data: Union[sc.DataArray, sc.Dataset], origin: str, target: str, scatter: bool
) -> Dict[Union[str, Tuple[str]], Callable]:
    """
    Get the conversion graph used by :py:func:`scippneutron.convert`
    when called with identical arguments.

    :param data: Input data.
    :param origin: Name of the input coordinate.
    :param target: Name of the output coordinate.
    :param scatter: Choose whether to use scattering or non-scattering conversions.
    :return: Conversion graph.
    :seealso: :py:func:`scippneutron.convert`, :py:func:`scippneutron.conversion_graph`.
    """
    return conversion_graph(
        origin, target, scatter, _deduce_energy_mode(data, origin, target)
    )


def convert(
    data: Union[sc.DataArray, sc.Dataset], origin: str, target: str, scatter: bool
) -> Union[sc.DataArray, sc.Dataset]:
    """
    Perform a unit conversion from the given origin unit to target.
    See the documentation page on "Coordinate Transformations"
    (https://scipp.github.io/scippneutron/user-guide/coordinate-transformations.html)
    for more details.

    :param data: Input data.
    :param origin: Name of the input coordinate.
    :param target: Name of the output coordinate.
    :param scatter: Choose whether to use scattering or non-scattering conversions.
    :return: A new scipp.DataArray or scipp.Dataset with the new coordinate.
    :seealso: :py:func:`scippneutron.deduce_conversion_graph` and
              :py:func:`scippneutron.conversion_graph` to inspect
              the possible conversions.
    """

    graph = deduce_conversion_graph(data, origin, target, scatter)

    try:
        converted = data.transform_coords(target, graph=graph)
    except KeyError as err:
        if err.args[0] == target:
            raise RuntimeError(
                f"No viable conversion from '{origin}' to '{target}' "
                f"with scatter={scatter}."
            )
        raise RuntimeError(
            f"Missing coordinate '{err.args[0]}' for conversion "
            f"from '{origin}' to '{target}'"
        ) from None

    return converted
