# -*- coding: utf-8 -*-
#
# This file is part of the bliss project
#
# Copyright (c) 2015-2023 Beamline Control Unit, ESRF
# Distributed under the GNU LGPLv3. See LICENSE for more info.

"""
Contains helper to manage the `scan_info` metadata provided inside each scans.

The `scan_info` is dictionary compounded by basic python types to provides
metadata relative to a scan.

Most of the fields are automatically generated by BLISS.

But the dictionary can be feed with extra information, which can be normalized
or not. For normalized data, the class `ScanInfo` provides few helpers to
update them.

Acquisition chain
=================

The `scan_info` describe the acquisition chain of the scan.

It contains:

- Top master device description
    - An independent acquisition chains
    - It contains the a set of involved devices
- Device description
    - An acquisition object part of the chain
    - List of sub devices triggered (if some)
    - List of the channels generated
- Channel description
    - Channel for data generated by the scan
    - Provides channel metadata

Here is a representation of the `acquisition_chain` structure:


.. code-block:: python

    scan_info = {
        "acquisition_chain": {
            TOP_MASTER_NAME_1: {
                "devices": [DEVICE_MASTER_ID, DEVICE_SLAVE_ID, ...]
            },
            ...
        },
        "devices": {
            DEVICE_MASTER_ID: {
                "channels": [CHANNEL_PATH_1, CHANNEL_PATH_2, ...],
                "triggered_devices": [DEVICE_SLAVE_ID, ...]
                ...
            },
            DEVICE_SLAVE_ID: {
                "channels": [CHANNEL_PATH_3, CHANNEL_PATH_4, ...]
                ...
            },
        },
        "channels": {
            CHANNEL_PATH_1: {...},
            CHANNEL_PATH_2: {...},
            CHANNEL_PATH_3: {...},
            CHANNEL_PATH_4: {...},
            ...
        },
    }

You better not to do assumption on the we way the devices/channels keys are
generated. But for now that's the only way to reach there names.

For now:

- `DEVICE_ID` are generated by `UNIQUE_ID + ":" + device_name`
    - Example: `31614654:timer`
- `CHANNEL_PATH` are generated usign the parent names (excuding the top master) separated by ":".
    - Example: `timer:elapsed_time`

So reaching name for devices or channels can be done the following way:
`name = key.rsplit(":", 1)[-1]`.

Device metadata
===============

Metadata can be exposed per devices. It can be feed by the controller or by the
acquisition object during the preparation of the controllers.

BLISS provides a type metadata only 2 kind of objects: `mca` and `lima`

====================  ===========  ===========  ========================
 Name                  Type         Flag         Description
====================  ===========  ===========  ========================
`channels`             list[str]    optional     If set, list of channels exposed
                                                 by this device
`triggered_devices`    list[str]    optional     If set, list of sub devices
                                                 triggered by this device
`type`                 str          optional     One of `lima` or `mca`
                                                 (other values could be used
                                                 but are not yet normalized)
`xaxis_channel`        str          optional     If defined, 1D channels exposed
                                                 by this device should be displayed
                                                 using this channel as x-axis
`xaxis_array`          ndarray      optional     If defined, 1D channels exposed
                                                 by this device should be displayed
                                                 using this array as x-axis
`xaxis_array_label`    str          optional     Label for the x-axis if `xaxis_array`
                                                 is used
`xaxis_array_unit`     str          optional     Unit for the x-axis if `xaxis_array`
                                                 is used
====================  ===========  ===========  ========================

This can be used to infer sub devices and channels.

- A device typed as `mca`
   - contains 1D channels which are spectrums
- A device typed as `lima`
   - can contain an 2D channel which is the image of the detector
   - can contain a device named `bpm`
   - can contain a device named `roi_counters`
      - this device contains a bunch of 0D channels associated to a BMP
   - can contain a device named `roi_profiles`
      - this device contains a bunch of 1D channels associated to ROIs
   - can contain a device named `roi_collection`
      - this device contains a bunch of 1D channels, each value is associated to a single ROI

Channel metadata
================

Each channel metadata are stored per channel names in the filed `channels` of
the scan_info.

The following metadata are automatically generated.

==============  ======  ===========  ========================
 Name            Type    Flag          Description
==============  ======  ===========  ========================
`display_name`   str     optional      Expected displayed name
`unit`           str     optional      SI unit used by the data
`dim`            int     mandatory     Dimensionality of the data
                                       (0: scalar, 1: 1D data, 2: 2D data)
==============  ======  ===========  ========================

Other metadata can be feed by the scan designers.
Take a look at :meth:`ScanInfo.set_channel_meta`.


API
===
"""

import typing
from bliss.common.utils import typecheck
import numbers
import logging


_logger = logging.getLogger(__name__)


class ScanInfo(dict):
    """
    Holder of metadata associated to a scan.

    It provides a key-value API to store metadata plus helper to feed this
    dictionary.

    It is exposed as a normal dictionary by Redis scan nodes.
    """

    def __init__(self):
        self._scan_info = self

    def _set_scan_info(self, scan_info):
        """This method and attribute is a trick to make the backward compatible
        implementation of ScanInfoFactory easy. This have to be removed together
        with ScanInfoFactory in few version.

        Introduced with BLISS 1.7
        """
        self._scan_info = scan_info

    @staticmethod
    def normalize(scan_info):
        """Returns a ScanInfo initialized from a raw object"""
        if isinstance(scan_info, ScanInfo):
            return scan_info
        result = ScanInfo()
        if scan_info is None:
            pass
        elif isinstance(scan_info, dict):
            result.update(scan_info)
        else:
            assert False
        return result

    @typecheck
    def set_channel_meta(
        self,
        name: str,
        start: typing.Optional[numbers.Real] = None,
        stop: typing.Optional[numbers.Real] = None,
        min: typing.Optional[numbers.Real] = None,
        max: typing.Optional[numbers.Real] = None,
        points: typing.Optional[numbers.Integral] = None,
        axis_points: typing.Optional[numbers.Integral] = None,
        axis_kind: typing.Optional[str] = None,
        group: typing.Optional[str] = None,
        axis_id: typing.Optional[numbers.Integral] = None,
        axis_points_hint: typing.Optional[numbers.Integral] = None,
    ):
        """
        Define metadata relative to a channel name

        Arguments:
            name: Name of the channel
            start: Start position of the axis
            stop: Stop position of the axis
            min: Minimal value the channel can have
            max: Minimal value the channel can have
            points: Amount of total points which will be transmitted by this channel
            axis_points: Amount of points for the axis (see scatter below)
            axis_kind: Kind of axis. It is used to speed up solid rendering in
                GUI. Can be one of:

                - `forth`: Move from start to stop always
                - `backnforth`: Move from start to stop to start
                - `step`: The motor position is discrete. The value can be used\
                  to group data together.

            group: Specify a group for the channel. All the channels from the
                same group are supposed to contain the same amount of item at
                the end of the scan. It also can be used as a hint for
                interactive user selection.
            axis_id: Index of the axis in the scatter. 0 is the fastest.
            axis_points_hint: Number of approximate points expected in the axis
                when this number of points is not regular
        """
        channels = self._scan_info.setdefault("channels", {})
        assert axis_kind in set([None, "forth", "backnforth", "step"])
        meta = channels.setdefault(name, {})
        if start is not None:
            meta["start"] = start
        if stop is not None:
            meta["stop"] = stop
        if min is not None:
            meta["min"] = min
        if max is not None:
            meta["max"] = max
        if points is not None:
            meta["points"] = points
        if axis_points is not None:
            meta["axis-points"] = axis_points
        if axis_id is not None:
            assert axis_id >= 0
            meta["axis-id"] = axis_id
        if axis_kind is not None:
            meta["axis-kind"] = axis_kind
        if group is not None:
            meta["group"] = group
        if axis_points_hint is not None:
            meta["axis-points-hint"] = axis_points_hint

        if axis_kind is not None:
            # If a kind is set the orientation of the scan have to be known
            if start is None:
                _logger.warning(
                    "A 'start' position for the channel '%s' should be set", name
                )
            if stop is None:
                _logger.warning(
                    "A 'stop' position for the channel '%s' should be set", name
                )

    def add_plots_entry(self):
        """CReate the main entry for plots description"""
        self._scan_info.setdefault("plots", [])

    @typecheck
    def add_scatter_plot(
        self,
        name: typing.Optional[str] = None,
        x: typing.Optional[str] = None,
        y: typing.Optional[str] = None,
        value: typing.Optional[str] = None,
    ):
        """
        Add a scatter plot definition to this `scan_info`.

        This can be used as default plot for the scan.

        Arguments:
            name: Unique name for the plot. If not defined, it is considered as a
                "default" plot
            x: Channel name for the x-axis
            y: Channel name for the y-axis
            value: Channel name for the data value
        """
        plots = self._scan_info.setdefault("plots", [])
        if not isinstance(plots, list):
            raise TypeError("The 'plots' metadata is corrupted. A list is expected.")

        item = {"kind": "scatter"}
        if x is not None:
            item["x"] = x
        if y is not None:
            item["y"] = y
        if value is not None:
            item["value"] = value

        items = []
        if len(item) > 1:
            items.append(item)

        plot = {"kind": "scatter-plot", "items": items}
        if name is not None:
            plot["name"] = name

        plots.append(plot)

    def has_default_curve_plot(self) -> bool:
        """Returns true if a curve plot is already defined"""
        plots = self._scan_info.get("plots", [])
        for plot in plots:
            if plot["kind"] == "curve-plot":
                if plot.get("name") is None:
                    return True
        return False

    @typecheck
    def add_curve_plot(
        self,
        name: typing.Optional[str] = None,
        x: typing.Optional[str] = None,
        yleft: typing.Union[None, typing.List[str], str] = None,
        yright: typing.Union[None, typing.List[str], str] = None,
    ):
        """
        Add a curve plot definition to this `scan_info`.

        It is an helper to simplify the creation of many curves.

        This can replace the default plot for the scan.

        Arguments:
            name: Unique name for the plot. If not defined, it is considered as a
                "default" plot
            x: Channel name for the x-axis
            yleft: Channel names of the curves which have to be displayed in the left y-axis
            yright: Channel names of the curves which have to be displayed in the right y-axis
        """
        plots = self._scan_info.setdefault("plots", [])
        if not isinstance(plots, list):
            raise TypeError("The 'plots' metadata is corrupted. A list is expected.")

        items = []
        if x is not None and yleft is None and yright is None:
            item = {"kind": "curve", "x": x}
            items.append(item)
        if yleft is not None:
            if isinstance(yleft, str):
                yleft = [yleft]
            for y in yleft:
                item = {"kind": "curve", "y_axis": "left"}
                if x is not None:
                    item["x"] = x
                item["y"] = y
                items.append(item)
        if yright is not None:
            if isinstance(yright, str):
                yright = [yright]
            for y in yright:
                item = {"kind": "curve", "y_axis": "right"}
                if x is not None:
                    item["x"] = x
                item["y"] = y
                items.append(item)

        plot = {"kind": "curve-plot", "items": items}
        if name is not None:
            plot["name"] = name

        plots.append(plot)

    @typecheck
    def add_1d_plot(
        self,
        name: typing.Optional[str] = None,
        x: typing.Optional[str] = None,
        y: typing.Union[None, typing.List[str], str] = None,
    ):
        """
        Add a 1D plot definition to this `scan_info`.

        It is an helper to simplify the creation of many curves.

        Arguments:
            name: Unique name for the plot. If not defined, it is considered as a
                "default" plot
            x: Channel name for the x-axis
            y: Channel names of the curves which have to be displayed in the y-axis
        """
        plots = self._scan_info.setdefault("plots", [])
        if not isinstance(plots, list):
            raise TypeError("The 'plots' metadata is corrupted. A list is expected.")

        items = []
        if y is not None:
            if isinstance(y, str):
                y = [y]
            for y in y:
                item = {"kind": "curve", "y": y}
                items.append(item)

        plot = {"kind": "1d-plot", "items": items}
        if x is not None:
            plot["x"] = x
        if name is not None:
            plot["name"] = name

        plots.append(plot)

    def set_sequence_info(self, scan_count: typing.Optional[int] = None):
        """
        Set extra-info for a sequence.

        Arguments:
            scan_count: Set it if you know the amount of scan which will be part
                        your sequence. THis can be used to know client side the
                        progress of the sequence.
        """
        info = self._scan_info.setdefault("sequence-info", {})
        info["scan-count"] = int(scan_count)

    def _get_key_from_acq_obj(self, acq_obj):
        """Used device key from an acq_obj"""
        return f"{str(id(acq_obj))}:{acq_obj.name}"

    def _set_device_meta(self, acq_obj, metadata: typing.Dict):
        """Protected function called by the scan to feed the device info after
        the device preparation.

        Argument:
            acq_obj: A device though its acquisition object
            metadata: A dictionary of basic python types.
                      Few reserved keys could be ignored.
        """
        devices = self.setdefault("devices", {})
        name = self._get_key_from_acq_obj(acq_obj)
        device = devices.setdefault(name, {})
        device_metadata = device.setdefault("metadata", {})
        device_metadata.update(metadata)

    def _get_channels_dict(self, acq_object, channels_dict):
        scalars = channels_dict.setdefault("scalars", [])
        spectra = channels_dict.setdefault("spectra", [])
        images = channels_dict.setdefault("images", [])
        already_read = set([])

        for acq_chan in acq_object.channels:
            fullname = acq_chan.fullname
            if fullname in already_read:
                continue
            already_read.add(fullname)
            shape = acq_chan.shape
            if len(shape) == 0 and fullname not in scalars:
                scalars.append(fullname)
            elif len(shape) == 1 and fullname not in spectra:
                spectra.append(fullname)
            elif len(shape) == 2 and fullname not in images:
                images.append(fullname)

        return channels_dict

    def _device_fullname(self, acq_obj):
        """Returns the full name of this device in order to make it unique."""
        return f"{id(acq_obj)}:{acq_obj.name}"

    def set_acquisition_chain_info(self, acq_chain):
        """
        Go through this acquisition chain, group acquisition channels by master
        and data shape.
        """
        from bliss.scanning.chain import AcquisitionMaster

        tree = acq_chain.tree

        chain_dict = {}
        for path in tree.paths_to_leaves():
            master = None
            # path[0] is root
            for acq_object in path[1:]:
                # it is mandatory to find an acq. master first
                if isinstance(acq_object, AcquisitionMaster):
                    if master is None or acq_object.parent is None:
                        master = acq_object.name
                        channels = chain_dict.setdefault(master, {"master": {}})
                        self._get_channels_dict(acq_object, channels["master"])
                        continue
                self._get_channels_dict(acq_object, channels)
        self._scan_info["acquisition_chain"] = chain_dict

        devices = self._scan_info.setdefault("devices", {})
        channels = self._scan_info.setdefault("channels", {})

        # Update devices keys
        for top_master_node in tree.children(tree.root):
            top_master = top_master_node.identifier
            if top_master is None:
                continue
            top_master_devices = chain_dict[top_master.name].setdefault("devices", [])
            for acq_object in tree.expand_tree(top_master):
                if acq_object is None:
                    continue
                acq_dev_name = self._device_fullname(acq_object)
                top_master_devices.append(acq_dev_name)
                device_info = devices.setdefault(acq_dev_name, {})
                triggered_devices = [
                    self._device_fullname(d.identifier)
                    for d in tree.children(acq_object)
                ]
                if triggered_devices:
                    device_info["triggered_devices"] = triggered_devices
                channel_names = []
                for acq_chan in acq_object.channels:
                    channel_names.append(acq_chan.fullname)
                    meta = channels.setdefault(acq_chan.fullname, {})
                    meta["device"] = acq_dev_name
                device_info["channels"] = channel_names

        # Update channels key
        for path in tree.paths_to_leaves():
            for acq_object in path[1:]:
                for acq_chan in acq_object.channels:
                    fullname = acq_chan.fullname
                    meta = channels.setdefault(fullname, {})
                    unit = acq_chan.unit
                    if unit is not None:
                        meta["unit"] = unit
                    display_name = acq_chan.short_name
                    if display_name is not None:
                        meta["display_name"] = display_name
                    shape = acq_chan.shape
                    if shape is not None:
                        meta["dim"] = len(shape)
