# -*- 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.

import numpy
from collections import deque

from bliss.scanning.chain import AcquisitionSlave, ChainNode
from bliss.scanning.channel import AcquisitionChannel
from bliss.common.event import dispatcher


class CalcHook(object):
    def compute(self, sender, data_dict):
        raise NotImplementedError

    def prepare(self):
        pass

    def start(self):
        pass

    def stop(self):
        pass


class CalcAcquisitionSlaveBase(AcquisitionSlave):
    """Emits data based on data emitted by other acquisition objects."""

    def __init__(self, *args, **kwargs):
        self._connected = False
        super().__init__(*args, **kwargs)

    def prepare(self):
        self.connect()

    def start(self):
        pass

    def stop(self):
        self.disconnect()

    def connect(self):
        if self._connected:
            return
        for channel in self._iter_input_channels():
            if channel.reference:
                dispatcher.connect(self.new_data_received, "new_data_stored", channel)
            else:
                dispatcher.connect(self.new_data_received, "new_data", channel)
        self._connected = True

    def disconnect(self):
        if not self._connected:
            return
        for channel in self._iter_input_channels():
            if channel.reference:
                dispatcher.disconnect(
                    self.new_data_received, "new_data_stored", channel
                )
            else:
                dispatcher.disconnect(self.new_data_received, "new_data", channel)
        self._connected = False

    def _iter_input_channels(self):
        raise NotImplementedError

    def _iter_output_channels(self):
        yield from self.channels

    def new_data_received(self, event_dict=None, signal=None, sender=None):
        input_data = self._prepare_calc_input(event_dict, signal, sender)
        if input_data is not None:
            calc_data = self.compute(sender, input_data)
            self._emit_calc_data(calc_data)

    def _prepare_calc_input(self, event_dict, signal, sender):
        input_data = event_dict.get("data")
        if input_data is not None and len(input_data) == 0:
            return None
        return input_data

    def compute(self, sender, input_data: dict) -> dict:
        raise NotImplementedError

    def _emit_calc_data(self, calc_data: dict):
        if not calc_data:
            return
        for channel in self._iter_output_channels():
            # Get channel data from calculation
            if self.device is None:
                # note: use channel **name** (see issue #3026)
                output_data = calc_data.get(channel.name)
            else:
                output_data = calc_data.get(self.device.tags[channel.short_name])
            if output_data is None:
                continue
            output_data = numpy.asarray(output_data)
            ndim_detector = output_data.ndim - 1

            # Set channel's data properties
            channel.shape = output_data.shape[1:]
            channel.dtype = output_data.dtype

            # Set Redis node's data properties
            node = channel.data_node
            if node.dtype is None:
                node.info["dtype"] = output_data.dtype
            if node.shape is None or len(node.shape) != ndim_detector:
                node.info["shape"] = (0,) * ndim_detector

            channel.emit(output_data)


class CalcChannelAcquisitionSlave(CalcAcquisitionSlaveBase):
    """
    Helper to do some extra Calculation on channels.
    i.e: compute encoder position to user position
    Args:
        src_acq_devices_list -- list or tuple of acq(device/master) you want to listen to.
        func -- the transformation function. This will have has input a  dictionary
        with the name of counter as the key and the value has the data of source data channel.
        This function should return a dictionary with the name of the destination channel as key,
        and the value as its data.
        Can also be an inherited class of **CalcHook**:
         - the transformation function is the **compute** method.
         - optionally you can redefine prepare,start,stop.
    """

    def __init__(
        self,
        name,
        src_acq_devices_list,
        func,
        output_channels_list,
        prepare_once=False,
        start_once=False,
    ):
        super().__init__(
            None,
            name=name,
            trigger_type=AcquisitionSlave.HARDWARE,
            prepare_once=prepare_once,
            start_once=start_once,
        )

        self.src_acq_devices_list = src_acq_devices_list

        if isinstance(func, CalcHook):
            self.cbk = func
        else:

            class CBK(CalcHook):
                def compute(self, sender, data_dict):
                    return func(sender, data_dict)

            self.cbk = CBK()

        for chan_out in output_channels_list:
            if isinstance(chan_out, AcquisitionChannel):
                self.channels.append(chan_out)
            elif isinstance(chan_out, str):
                self.channels.append(AcquisitionChannel(chan_out, float, ()))
            else:
                raise TypeError(f"Object '{chan_out}'' is not an AcquisitionChannel")

    def _iter_input_channels(self):
        for acq_device in self.src_acq_devices_list:
            for channel in acq_device.channels:
                yield channel

    def prepare(self):
        self.cbk.prepare()
        super().prepare()

    def start(self):
        self.cbk.start()
        super().start()

    def stop(self):
        super().stop()
        self.cbk.stop()

    def _prepare_calc_input(self, event_dict, signal, sender):
        input_data = super()._prepare_calc_input(event_dict, signal, sender)
        return {sender.short_name: input_data}

    def compute(self, sender, input_data: dict) -> dict:
        return self.cbk.compute(sender, input_data)


class CalcCounterAcquisitionSlave(CalcAcquisitionSlaveBase):
    """
    Helper to do some extra Calculation on counters.
    i.e: compute encoder position to user position
    Args:
        controller -- CalcCounterController Object
        src_acq_devices_list -- list or tuple of acq(device/master) you want to listen to.
    """

    def __init__(self, controller, src_acq_devices_list, acq_params, ctrl_params=None):
        super().__init__(
            controller,
            name=controller.name,
            npoints=acq_params.get("npoints", 1),
            trigger_type=AcquisitionSlave.HARDWARE,
            ctrl_params=ctrl_params,
        )

        self._frame_index = dict()
        self._inputs_channels = dict()
        self.build_input_channel_list(src_acq_devices_list)

    def build_input_channel_list(self, src_acq_devices_list):
        for acq_device in src_acq_devices_list:
            for cnt, channels in acq_device._counters.items():
                # filter unwanted counters and extra channels
                if cnt in self.device._input_counters:
                    # ignore multi channels per counter (see sampling)
                    self._inputs_channels[channels[0]] = cnt

        self._inputs_data_buffer = {chan: deque() for chan in self._inputs_channels}

    def _iter_input_channels(self):
        yield from self._inputs_channels

    def _prepare_calc_input(self, event_dict, signal, sender):
        if event_dict["description"]["reference"]:
            # Handle Lima image reference
            if event_dict["data"].get("in_prepare"):
                return

            curidx = self._frame_index.setdefault(sender, 0)
            input_data = sender.data_node.get_as_array(curidx, -1)

            # if there is no data to grab
            if input_data.size == 0:
                return

            # case of a single frame without a stacking dimension (only happen with references)
            if input_data.ndim == 2:
                input_data = input_data[numpy.newaxis, ...]

            self._frame_index[sender] += input_data.shape[0]
            return input_data
        return super()._prepare_calc_input(event_dict, signal, sender)

    def compute(self, sender, sender_data) -> dict:
        """
        This method works only if all input_counters will generate the same number of points !!!
        It registers all data comming from the input counters.
        It calls calc_function with input counters data which have reach the same index
        This function is called once per counter (input and output).

        * <sender> = AcquisitionChannel
        * <data_dict> = {'em1ch1': array([0.00256367])}
        """

        # buffering: tmp storage of received newdata
        self._inputs_data_buffer[sender].extend(sender_data)

        # Find the amount of aligned data (i.e the smallest newdata len among all inputs)
        # Build the input_data_dict (indexed by tags and containing aligned data for all inputs)
        # Pop data from _inputs_data_buffer while building input_data_dict

        aligned_data_index = min(
            [len(data) for data in self._inputs_data_buffer.values()]
        )
        if aligned_data_index > 0:
            input_data_dict = dict()
            for chan, cnt in self._inputs_channels.items():
                aligned_data = [
                    self._inputs_data_buffer[chan].popleft()
                    for i in range(aligned_data_index)
                ]
                input_data_dict[self.device.tags[cnt.name]] = numpy.array(aligned_data)

            output_data_dict = self.device.calc_function(input_data_dict)

            return output_data_dict


class CalcCounterChainNode(ChainNode):
    def get_acquisition_object(
        self, acq_params, ctrl_params=None, parent_acq_params=None
    ):

        # Check if Acquisition Devices of dependant counters already exist
        acq_devices = []
        for node in self._calc_dep_nodes.values():
            acq_obj = node.acquisition_obj
            if acq_obj is None:
                raise ValueError(
                    f"cannot create CalcCounterAcquisitionSlave: acquisition object of {node}({node.controller}) is None!"
                )
            else:
                acq_devices.append(acq_obj)

        return self.controller.get_acquisition_object(
            acq_params, ctrl_params, parent_acq_params, acq_devices
        )
