import warnings

import torch
from numpy import product

from sinabs.layers import NeuromorphicReLU


def synops_hook(layer, inp, out):
    assert len(inp) == 1, "Multiple inputs not supported for synops hook"
    inp = inp[0]
    layer.tot_in = inp.sum()
    layer.tot_out = out.sum()
    layer.synops = layer.tot_in * layer.fanout
    layer.tw = inp.shape[0]


class SNNSynOpCounter:
    """
    Counter for the synaptic operations emitted by all SpikingLayers in a
    spiking model.
    Note that this is automatically instantiated by `from_torch` and by
    `Network` if they are passed `synops=True`.

    Arguments:
        model: Spiking model.
        dt: the number of milliseconds corresponding to a time step in the \
        simulation (default 1.0).

    Example:
        >>> counter = SNNSynOpCounter(my_spiking_model)
        >>> output = my_spiking_model(input)  # forward pass
        >>> synops_table = counter.get_synops()
    """

    def __init__(self, model, dt=1.0):
        self.model = model
        self.handles = []
        self.dt = dt

        for layer in model.modules():
            self._register_synops_hook(layer)

    def _register_synops_hook(self, layer):
        if isinstance(layer, torch.nn.Conv2d):
            layer.fanout = (
                layer.out_channels * product(layer.kernel_size) / product(layer.stride)
            )
        elif isinstance(layer, torch.nn.Linear):
            layer.fanout = layer.out_features

        else:
            return None

        handle = layer.register_forward_hook(synops_hook)
        self.handles.append(handle)

    def get_synops(self) -> dict:
        """Method to compute a table of synaptic operations for the latest forward pass.

        .. note:: this may not be accurate in presence of average pooling.

        Returns:
            A dictionary containing layer IDs and respectively, for
            the latest forward pass performed, their number of input spikes,
            fanout, synaptic operations, number of timesteps, total duration
            of simulation, number of synaptic operations per second.

        Example:
            >>> synops_map = counter.get_synops()
            >>> SynOps_dataframe = pandas.DataFrame.from_dict(synops_map, "index")
            >>> SynOps_dataframe.set_index("Layer", inplace=True)

        """
        SynOps_map = {}
        scale_facts = []
        for i, lyr in enumerate(self.model.modules()):
            if isinstance(lyr, torch.nn.AvgPool2d):
                if lyr.kernel_size != lyr.stride:
                    warnings.warn(
                        f"In order for the Synops counter to work accurately the pooling "
                        f"layers kernel size should match their strides. At the moment at layer {i}, "
                        f"the kernel_size = {lyr.kernel_size}, the stride = {lyr.stride}."
                    )
                ks = lyr.kernel_size
                scale_factor = ks**2 if isinstance(ks, int) else ks[0] * ks[1]
                scale_facts.append(scale_factor)
            if hasattr(lyr, "synops"):
                scale_factor = 1
                while len(scale_facts) != 0:
                    scale_factor *= scale_facts.pop()
                SynOps_map[i] = {
                    "Layer": i,
                    "In": lyr.tot_in * scale_factor,
                    "Fanout_Prev": lyr.fanout,
                    "SynOps": lyr.synops * scale_factor,
                    "N. timesteps": lyr.tw,
                    "Time window (ms)": lyr.tw * self.dt,
                    "SynOps/s": (lyr.synops * scale_factor) / lyr.tw / self.dt * 1000,
                }

        return SynOps_map

    def get_total_synops(self, per_second=False) -> float:
        """
        Sums up total number of synaptic operations across the network.

        .. note:: this may not be accurate in presence of average pooling.

        Arguments:
            per_second (bool, default False): if True, gives synops per second
                            instead of total synops in the last forward pass.

        Returns:
            synops: the total synops in the network, based on the last forward pass.
        """
        synops = 0.0
        for i, lyr in enumerate(self.model.modules()):
            if hasattr(lyr, "synops"):
                if per_second:
                    layer_synops = lyr.synops / lyr.tw / self.dt * 1000
                else:
                    layer_synops = lyr.synops

                synops = synops + layer_synops
        return synops

    def get_total_power_use(self, j_per_synop=1e-11):
        """
        Method to quickly get the total power use of the network, estimated
        over the latest forward pass.

        Arguments:
            j_per_synop: Energy use per synaptic operation, in joules.\
            Default 1e-11 J.

        Returns: 
            estimated power in mW.
        """
        tot_synops_per_s = self.get_total_synops(per_second=True)
        power_in_mW = tot_synops_per_s * j_per_synop * 1000
        return power_in_mW

    def __del__(self):
        for handle in self.handles:
            handle.remove()


class SynOpCounter:
    """
    Counter for the synaptic operations emitted by all Neuromorphic ReLUs in an
    analog CNN model.

    Parameters:
        modules: list of modules, e.g. MyTorchModel.modules()
        sum_activations: If True (default), returns a single number of synops, otherwise a list of layer synops.

    Example:
        >>> counter = SynOpCounter(MyTorchModel.modules(), sum_activations=True)
        >>> output = MyTorchModule(input)  # forward pass
        >>> synop_count = counter()
    """

    def __init__(self, modules, sum_activations=True):
        self.modules = []
        for module in modules:
            if isinstance(module, NeuromorphicReLU) and module.fanout > 0:
                self.modules.append(module)

        if len(self.modules) == 0:
            raise ValueError("No NeuromorphicReLU found in module list.")

        self.sum_activations = sum_activations
        # self.modules[1:] = []

    def __call__(self):
        synops = []
        for module in self.modules:
            synops.append(module.activity)

        if self.sum_activations:
            synops = torch.stack(synops).sum()
        return synops
