from pymonntorch import NetworkObject, Network, NeuronDimension, Behavior, NeuronGroup
from conex.behaviors.neurons.axon import NeuronAxon
from conex.behaviors.neurons.setters import SensorySetter, LocationSetter
from conex.behaviors.neurons.specs import SpikeTrace
from conex.behaviors.neurons.dendrite import SimpleDendriteStructure
from torch.utils.data.dataloader import DataLoader
from typing import Union, Dict, Callable, Tuple, List
from conex.nn.priority import LAYER_PRIORITIES, NEURON_PRIORITIES
from conex.behaviors.layer.dataset import SpikeNdDataset
from conex.nn.structure.port import Port


class InputLayer(NetworkObject):
    """A sample input layer.

    Args:
        input_dataloader (torch dataloader): The dataloader of input layer.
        have_sensory (bool): Whether dataloader return sensory data or not.
        have_location (bool): Whether dataloader location data or not.
        have_label (bool): Whether dataloader label or not.
        sensory_size (int or behavior): The size of each sensory neurongroup.
        location_size (int or behavior): The size of each location neurongroup.
        sensory_axon (behavior): The behavior class for axon paradigm of sensory neurongroup.
        sensory_axon_params (dict): Parameters for axon class of sensory neurongroup.
        location_axon (behavior): The behavior class for axon paradigm of location neurongroup.
        location_axon_params (dict): Parameters for axon class of location neurongroup.
        silent_interval (int): Empty interval between two samples.
        instance_duration (int): Each sample duraiton
        sensory_trace (float): The spike trace of sensory neurongroup.
        location_trace (float): The spike trace of location neurongroup.
        sensory_data_dim (int): The number of dimension of sensory data.
        location_data_dim (int): The number of dimension of location data.
        behavior (dict): The behavior for the InputLayer itself.
        tag (str): tag of the InputLayer divided by ",".
        sensory_tag (str): tag of the sensory population divided by ",".
        location_tag (ste): tag of the location population divided by ",".
        sensory_user_defined (dict): Additional behavior for sensory neurongroup.
        location_user_defined (dict): Additional behavior for location neurongroup.
    """

    def __init__(
        self,
        net: Network,
        input_dataloader: DataLoader,
        have_sensory: bool = True,
        have_location: bool = False,
        have_label: bool = True,
        sensory_size: Union[int, NeuronDimension] = None,
        location_size: Union[int, NeuronDimension] = None,
        sensory_axon: Behavior = NeuronAxon,
        sensory_axon_params: dict = None,
        location_axon: Behavior = NeuronAxon,
        location_axon_params: dict = None,
        silent_interval: int = 0,
        instance_duration: int = 0,
        sensory_trace: float = None,
        location_trace: float = None,
        sensory_data_dim: int = 2,
        location_data_dim: int = 2,
        tag: str = None,
        behavior: Dict[int, Behavior] = None,
        sensory_tag: str = None,
        location_tag: str = None,
        output_ports: Dict[
            str, Tuple[Union[Tuple, None], List[Tuple[str, Dict[int, Behavior]]]]
        ] = None,
        sensory_user_defined: Dict[int, Behavior] = None,
        location_user_defined: Dict[int, Behavior] = None,
    ):
        behavior = {} if behavior is None else behavior

        if LAYER_PRIORITIES["SpikeNdDataset"] not in behavior:
            behavior[LAYER_PRIORITIES["SpikeNdDataset"]] = SpikeNdDataset(
                dataloader=input_dataloader,
                ndim_sensory=sensory_data_dim,
                ndim_location=location_data_dim,
                have_location=have_location,
                have_sensory=have_sensory,
                have_label=have_label,
                silent_interval=silent_interval,
                instance_duration=instance_duration,
            )

        super().__init__(tag=tag, network=net, behavior=behavior, device=net.device)
        self.add_tag(self.__class__.__name__)

        self.network = net
        net.input_layers.append(self)

        sensory_tag = "Sensory" if sensory_tag is None else "Sensory," + sensory_tag

        if have_sensory and sensory_size is not None:
            self.sensory_pop = self.__get_ng(
                net=net,
                size=sensory_size,
                tag=sensory_tag,
                trace=sensory_trace,
                setter=SensorySetter,
                axon=sensory_axon,
                axon_params=sensory_axon_params,
                user_defined=sensory_user_defined,
            )

            self.sensory_pop.layer = self
            self.add_sub_structure(self.sensory_pop)

        location_tag = (
            "Location" if location_tag is None else "Location," + location_tag
        )

        if have_location and location_size is not None:
            self.location_pop = self.__get_ng(
                net=net,
                size=location_size,
                tag=location_tag,
                trace=location_trace,
                setter=LocationSetter,
                axon=location_axon,
                axon_params=location_axon_params,
                user_defined=location_user_defined,
            )

            self.location_pop.layer = self
            self.add_sub_structure(self.location_pop)

        self.output_ports = _create_port(output_ports, self)
        self.input_ports = {}

    def __get_ng(
        self,
        net: Network,
        size: Union[int, NeuronDimension],
        tag: Union[str, None],
        trace: Union[float, None],
        setter: Callable,
        axon: Behavior = NeuronAxon,
        axon_params: dict = None,
        user_defined: Dict[int, Behavior] = None,
    ) -> NeuronGroup:
        behavior = {
            NEURON_PRIORITIES["Fire"]: setter(),
        }

        if axon:
            params = axon_params if axon_params is not None else {}
            behavior[NEURON_PRIORITIES["NeuronAxon"]] = axon(**params)

        if trace is not None:
            behavior[NEURON_PRIORITIES["SpikeTrace"]] = SpikeTrace(tau_s=trace)

        if user_defined is not None:
            behavior.update(user_defined)

        return NeuronGroup(size=size, behavior=behavior, net=net, tag=tag)

    def __repr__(self) -> str:
        result = (
            self.__class__.__name__
            + "("
            + f"Sensory Population {self.sensory_pop.tags[0]}({self.sensory_pop.size})"
            if hasattr(self, "sensory_pop")
            else ""
            + f"Location Population {self.location_pop.tags[0]}({self.location_pop.size})"
            if hasattr(self, "location_pop")
            else "" + "){"
        )
        for k in sorted(list(self.behavior.keys())):
            result += str(k) + ":" + str(self.behavior[k])
        return result + "}"


class OutputLayer(NetworkObject):
    """A sample output layer.

    Args:
        representation_size (int or behavior): The size of each representation neurongroup.
        motor_size (int or behavior): The size of each motor neurongroup.
        representation_trace (float): The spike trace of representation neurongroup.
        motor_trace (float): The spike trace of motor neurongroup.
        representation_dendrite_structure (Callable): Dendrite structure for representation population.
        representation_dendrite_structure_params (dict): The parameters for dendrite structure of representation population.
        motor_dendrite_structure (Callable): Dendrite structure for motor population.
        motor_dendrite_structure_params (dict): The parameters for dendrite structure of representation population.
        behavior (dict): The behavior for the InputLayer itself.
        tag (str): tag of the InputLayer divided by ",".
        representation_tag (str): tag of the representation population divided by ",".
        motor_tag (str): tag of the motor population divided by ",".
        representation_user_defined (dict): Additional behavior for representation neurongroup.
        motor_user_defined (dict): Additional behavior for motor neurongroup.
    """

    def __init__(
        self,
        net: Network,
        representation_size: Union[int, NeuronDimension] = None,
        motor_size: Union[int, NeuronDimension] = None,
        representation_trace: Union[float, None] = None,
        motor_trace: Union[float, None] = None,
        representation_dendrite_structure: Callable = SimpleDendriteStructure,
        representation_dendrite_structure_params: dict = None,
        motor_dendrite_structure: Callable = SimpleDendriteStructure,
        motor_dendrite_structure_params: dict = None,
        tag: str = None,
        behavior: Dict[int, Behavior] = None,
        representation_tag: str = None,
        motor_tag: str = None,
        input_ports: Dict[
            str, Tuple[Union[Tuple, None], List[Tuple[str, Dict[int, Behavior]]]]
        ] = None,
        representation_user_defined: Dict[int, Behavior] = None,
        motor_user_defined: Dict[int, Behavior] = None,
    ):
        behavior = {} if behavior is None else behavior
        super().__init__(tag=tag, network=net, behavior=behavior, device=net.device)
        self.network = net
        net.output_layers.append(self)

        representation_tag = (
            "Representation"
            if representation_tag is None
            else "Representation," + representation_tag
        )

        if representation_size is not None:
            self.representation_pop = self.__get_ng(
                net=net,
                size=representation_size,
                tag=representation_tag,
                trace=representation_trace,
                dendrite_structure=representation_dendrite_structure,
                dendrite_structure_params=representation_dendrite_structure_params,
                user_defined=representation_user_defined,
            )

            self.representation_pop.layer = self
            self.add_sub_structure(self.representation_pop)

        motor_tag = "Motor" if motor_tag is None else "Motor," + motor_tag

        if motor_size is not None:
            self.motor_pop = self.__get_ng(
                net=net,
                size=motor_size,
                tag=motor_tag,
                trace=motor_trace,
                dendrite_structure=motor_dendrite_structure,
                dendrite_structure_params=motor_dendrite_structure_params,
                user_defined=motor_user_defined,
            )

            self.motor_pop.layer = self
            self.add_sub_structure(self.motor_pop)

        self.add_tag(self.__class__.__name__)

        self.input_ports = _create_port(input_ports, self)
        self.output_ports = {}

    def __get_ng(
        self,
        net: Network,
        size: Union[int, NeuronDimension],
        tag: Union[str, None],
        trace: Union[float, None] = None,
        dendrite_structure: Callable = SimpleDendriteStructure,
        dendrite_structure_params: dict = None,
        user_defined: Dict[int, Behavior] = None,
    ) -> NeuronGroup:
        dendrite_structure_params = (
            dendrite_structure_params if dendrite_structure_params is not None else {}
        )
        behavior = {}

        if dendrite_structure:
            behavior[NEURON_PRIORITIES["SimpleDendriteStructure"]] = dendrite_structure(
                **dendrite_structure_params
            )

        if trace is not None:
            behavior[NEURON_PRIORITIES["Trace"]] = SpikeTrace(tau_s=trace)

        if user_defined is not None:
            behavior.update(user_defined)

        return NeuronGroup(size, behavior, net, tag)

    def __repr__(self) -> str:
        result = (
            self.__class__.__name__
            + "("
            + f"representation Population {self.representation_pop.tags[0]}({self.representation_pop.size})"
            if hasattr(self, "representation_pop")
            else ""
            + f"Motor Population {self.motor_pop.tags[0]}({self.motor_pop.size})"
            if hasattr(self, "motor_pop")
            else "" + "){"
        )
        for k in sorted(list(self.behavior.keys())):
            result += str(k) + ":" + str(self.behavior[k])
        return result + "}"


def _create_port(
    port_dict: Dict[
        str, Tuple[Union[Tuple, None], List[Tuple[str, Dict[int, Behavior]]]]
    ],
    obj: Union[InputLayer, OutputLayer],
):
    result = {}
    if port_dict is not None:
        result = {
            k: (v[0], [Port(object=getattr(obj, sp[0]), behavior=sp[1]) for sp in v[1]])
            for k, v in port_dict.items()
        }
    return result
