import warnings

import numpy as np

from nengo.config import Config
from nengo.connection import Connection
from nengo.dists import Choice, Exponential, Uniform
from nengo.ensemble import Ensemble
from nengo.exceptions import ValidationError
from nengo.network import Network
from nengo.node import Node
from nengo.utils.network import with_self
from nengo.utils.numpy import is_iterable

from .ensemblearray import EnsembleArray


def filtered_step(t, shift=0.5, scale=50, step_val=1):
    return np.maximum(-1 / np.exp((t - shift) * scale) + 1, 0) * step_val


class AssociativeMemory(Network):
    """Associative memory network.

    Parameters
    ----------
    input_vectors: array_like
        The list of vectors to be compared against.
    output_vectors: array_like, optional
        The list of vectors to be produced for each match. If None, the
        associative memory will be autoassociative (cleanup memory).
    n_neurons: int, optional
        The number of neurons for each of the ensemble (where each ensemble
        represents each item in the input_vectors list).
    threshold: float, optional
        The association activation threshold.
    input_scales: float or array_like, optional
        Scaling factor to apply on each of the input vectors. Note that it
        is possible to scale each vector independently.
    inhibitable: bool, optional
        Flag to indicate if the entire associative memory module is
        inhibitable (entire thing can be shut off). The input gain into
        the inhibitory connection is 1.5.
    label : str, optional
        A name for the ensemble. Used for debugging and visualization.
    seed : int, optional
        The seed used for random number generation.
    add_to_container : bool, optional
        Determines if the network will be added to the current container.
        If None, will be true if currently within a Network.
    """

    exp_scale = 0.15  # Scaling factor for exponential distribution
    n_eval_points = 5000

    def __init__(  # noqa: C901
        self,
        input_vectors,
        output_vectors=None,
        n_neurons=50,
        threshold=0.3,
        input_scales=1.0,
        inhibitable=False,
        label=None,
        seed=None,
        add_to_container=None,
    ):
        super().__init__(label, seed, add_to_container)

        # --- Put arguments in canonical form
        if output_vectors is None:
            # If output vocabulary is not specified, use input vector list
            # (i.e autoassociative memory)
            output_vectors = input_vectors
        if is_iterable(input_vectors):
            input_vectors = np.array(input_vectors, ndmin=2)
        if is_iterable(output_vectors):
            output_vectors = np.array(output_vectors, ndmin=2)

        if input_vectors.shape[0] == 0:
            raise ValidationError(
                "Number of input vectors cannot be 0.", attr="input_vectors", obj=self
            )
        elif input_vectors.shape[0] != output_vectors.shape[0]:
            # Fail if number of input items and number of output items don't
            # match
            raise ValidationError(
                "Number of input vectors does not match number of output "
                f"vectors. {input_vectors.shape[0]} != {output_vectors.shape[0]}",
                attr="input_vectors",
                obj=type(self),
            )

        # Handle possible different threshold / input_scale values for each
        # element in the associative memory
        if not is_iterable(threshold):
            threshold = threshold * np.ones(input_vectors.shape[0])
        else:
            threshold = np.array(threshold)

        # --- Check preconditions
        self.n_items = input_vectors.shape[0]
        if threshold.shape[0] != self.n_items:
            raise ValidationError(
                f"Number of threshold values ({threshold.shape[0]}) does not match "
                f"number of input vectors ({self.n_items}).",
                attr="threshold",
                obj=self,
            )

        # --- Set parameters
        self.out_conns = []  # Used in `add_threshold_to_output`
        # Used in `add_threshold_to_output`
        self.default_vector_inhibit_conns = []
        self.thresh_ens = None  # Will hold thresholded outputs
        self.is_wta = False
        self._inhib_scale = 1.5

        # -- Create the core network
        with self, self.am_ens_config:
            self.bias_node = Node(output=1)
            self.elem_input = Node(size_in=self.n_items, label="element input")
            self.elem_output = Node(size_in=self.n_items, label="element output")
            self.utilities = self.elem_output

            self.am_ensembles = []
            label_prefix = "" if label is None else label + "_"
            filt_scale = 15
            filt_step_func = lambda x: filtered_step(x, 0.0, scale=filt_scale)
            for i in range(self.n_items):
                e = Ensemble(n_neurons, 1, label=label_prefix + str(i))
                self.am_ensembles.append(e)

                # Connect input and output nodes
                Connection(self.bias_node, e, transform=-threshold[i])
                Connection(self.elem_input[i], e)
                Connection(e, self.elem_output[i], function=filt_step_func)

            if inhibitable:
                # Input node for inhibitory gating signal (if enabled)
                self.inhibit = Node(size_in=1, label="inhibit")
                Connection(
                    self.inhibit,
                    self.elem_input,
                    transform=-np.ones((self.n_items, 1)) * self._inhib_scale,
                )
                # Note: We can use a decoded connection here because all the
                # am_ensembles have [1] encoders
            else:
                self.inhibit = None

            self.thresh_bias = None
            self.thresholded_utilities = None

        self.add_input_mapping("input", input_vectors, input_scales)
        self.add_output_mapping("output", output_vectors)

    @property
    def am_ens_config(self):
        """(Config) Defaults for associative memory ensemble creation."""
        cfg = Config(Ensemble, Connection)
        cfg[Ensemble].update(
            {
                "radius": 1,
                "intercepts": Exponential(self.exp_scale, 0.0, 1.0),
                "encoders": Choice([[1]]),
                "eval_points": Uniform(0.0, 1.0),
                "n_eval_points": self.n_eval_points,
            }
        )
        cfg[Connection].synapse = None
        return cfg

    @property
    def default_ens_config(self):
        """(Config) Defaults for other ensemble creation."""
        cfg = Config(Ensemble)
        cfg[Ensemble].update(
            {
                "radius": 1,
                "intercepts": Exponential(self.exp_scale, 0.0, 1.0),
                "encoders": Choice([[1]]),
                "eval_points": Uniform(0.0, 1.0),
                "n_eval_points": self.n_eval_points,
            }
        )
        return cfg

    @property
    def thresh_ens_config(self):
        """(Config) Defaults for threshold ensemble creation."""
        cfg = Config(Ensemble)
        cfg[Ensemble].update(
            {
                "radius": 1,
                "intercepts": Uniform(0.5, 1.0),
                "encoders": Choice([[1]]),
                "eval_points": Uniform(0.75, 1.1),
                "n_eval_points": self.n_eval_points,
            }
        )
        return cfg

    @with_self
    def add_input_mapping(self, name, input_vectors, input_scales=1.0):
        """Adds a set of input vectors to the associative memory network.

        Creates a transform with the given input vectors between the
        a named input node and associative memory element input to enable the
        inputs to be mapped onto ensembles of the Associative Memory.

        Parameters
        ----------
        name: str
            Name to use for the input node. This name will be used as the name
            of the attribute for the associative memory network.
        input_vectors: array_like
            The list of vectors to be compared against.
        input_scales: float or array_like, optional
            Scaling factor to apply on each of the input vectors. Note that it
            is possible to scale each vector independently.
        """
        # --- Put arguments in canonical form
        n_vectors, d_vectors = input_vectors.shape
        if is_iterable(input_vectors):
            input_vectors = np.array(input_vectors, ndmin=2)
        if not is_iterable(input_scales):
            input_scales = input_scales * np.ones((1, n_vectors))
        else:
            input_scales = np.array(input_scales, ndmin=2)

        # --- Check some preconditions
        if input_scales.shape[1] != n_vectors:
            raise ValidationError(
                f"Number of input_scale values ({input_scales.shape[1]}) does not "
                f"match number of input vectors ({n_vectors}).",
                attr="input_scales",
            )
        if hasattr(self, name):
            raise ValidationError(
                f"Name '{name}' already exists as a node in the associative memory.",
                attr="name",
            )

        # --- Finally, make the input node and connect it
        in_node = Node(size_in=d_vectors, label=name)
        setattr(self, name, in_node)
        Connection(
            in_node,
            self.elem_input,
            synapse=None,
            transform=input_vectors * input_scales.T,
        )

    @with_self
    def add_output_mapping(self, name, output_vectors):
        """Adds another output to the associative memory network.

        Creates a transform with the given output vectors between the
        associative memory element output and a named output node to enable the
        selection of output vectors by the associative memory.

        Parameters
        ----------
        name: str
            Name to use for the output node. This name will be used as
            the name of the attribute for the associative memory network.
        output_vectors: array_like
            The list of vectors to be produced for each match.
        """
        # --- Put arguments in canonical form
        if is_iterable(output_vectors):
            output_vectors = np.array(output_vectors, ndmin=2)

        # --- Check preconditions
        if hasattr(self, name):
            raise ValidationError(
                f"Name '{name}' already exists as a node in the associative memory.",
                attr="name",
            )

        # --- Make the output node and connect it
        output = Node(size_in=output_vectors.shape[1], label=name)
        setattr(self, name, output)

        if self.thresh_ens is not None:
            c = Connection(
                self.thresh_ens.output, output, synapse=None, transform=output_vectors.T
            )
        else:
            c = Connection(
                self.elem_output, output, synapse=None, transform=output_vectors.T
            )
        self.out_conns.append(c)

    @with_self
    def add_default_output_vector(
        self,
        output_vector,
        output_name="output",
        n_neurons=50,
        min_activation_value=0.5,
    ):
        """Adds a default output vector to the associative memory network.

        The default output vector is chosen if the input matches none of
        the given input vectors.

        Parameters
        ----------
        output_vector: array_like
            The vector to be produced if the input value matches none of
            the vectors in the input vector list.
        output_name: str, optional
            The name of the input to which the default output vector
            should be applied.
        n_neurons: int, optional
            Number of neurons to use for the default output vector ensemble.
        min_activation_value: float, optional
            Minimum activation value (i.e. threshold) to use to disable
            the default output vector.
        """
        with self.default_ens_config:
            default_vector_ens = Ensemble(
                n_neurons, 1, label=f"Default {output_name} vector"
            )
            Connection(self.bias_node, default_vector_ens, synapse=None)

            tr = -(1.0 / min_activation_value) * np.ones((1, self.n_items))
            if self.thresh_ens is not None:
                c = Connection(self.thresh_ens.output, default_vector_ens, transform=tr)
            else:
                c = Connection(self.elem_output, default_vector_ens, transform=tr)

            # Add the output connection to the output connection list
            self.default_vector_inhibit_conns.append(c)

            # Make new output class attribute and connect to it
            output = getattr(self, output_name)
            Connection(
                default_vector_ens,
                output,
                transform=np.array(output_vector, ndmin=2).T,
                synapse=None,
            )

            if self.inhibit is not None:
                Connection(
                    self.inhibit,
                    default_vector_ens,
                    transform=-self._inhib_scale,
                    synapse=None,
                )

    @with_self
    def add_wta_network(self, inhibit_scale=1.5, inhibit_synapse=0.005):
        """Add a winner-take-all (WTA) network to associative memory output.

        Parameters
        ----------
        inhibit_scale: float, optional
            Mutual inhibition scaling factor.
        inhibit_synapse: float, optional
            Mutual inhibition synapse time constant.
        """
        if not self.is_wta:
            Connection(
                self.elem_output,
                self.elem_input,
                synapse=inhibit_synapse,
                transform=((np.eye(self.n_items) - 1) * inhibit_scale),
            )
            self.is_wta = True
        else:
            warnings.warn(
                "AssociativeMemory network is already configured "
                "with a WTA network. Additional `add_wta_network` "
                "calls are ignored."
            )

    @with_self
    def add_threshold_to_outputs(self, n_neurons=50, inhibit_scale=10):
        """Adds a thresholded output to the associative memory.

        Parameters
        ----------
        n_neurons: int, optional
            Number of neurons to use for the default output vector ensemble.
        inhibit_scale: float, optional
            Mutual inhibition scaling factor.
        """
        if self.thresh_ens is not None:
            warnings.warn(
                "AssociativeMemory network is already configured "
                "with thresholded outputs. Additional "
                "`add_threshold_to_output` calls are ignored."
            )
            return

        with self.thresh_ens_config:
            self.thresh_bias = EnsembleArray(
                n_neurons, self.n_items, label="thresh_bias"
            )
            self.thresh_ens = EnsembleArray(n_neurons, self.n_items, label="thresh_ens")
            self.thresholded_utilities = self.thresh_ens.output

            Connection(
                self.bias_node,
                self.thresh_bias.input,
                transform=np.ones((self.n_items, 1)),
                synapse=None,
            )
            Connection(
                self.bias_node,
                self.thresh_ens.input,
                transform=np.ones((self.n_items, 1)),
                synapse=None,
            )
            Connection(
                self.elem_output, self.thresh_bias.input, transform=-inhibit_scale
            )
            Connection(
                self.thresh_bias.output, self.thresh_ens.input, transform=-inhibit_scale
            )

            # Reroute connections from elem_output to be connections
            # from thresh_ens.output
            def reroute(conn):
                c = Connection(
                    self.thresh_ens.output,
                    conn.post,
                    transform=conn.transform,
                    synapse=conn.synapse,
                )
                self.connections.remove(conn)
                return c

            self.default_vector_inhibit_conns = [
                reroute(c) for c in self.default_vector_inhibit_conns
            ]
            self.out_conns = [reroute(c) for c in self.out_conns]

            # Make inhibitory connection if inhibit option is set
            if self.inhibit is not None:
                for e in self.thresh_ens.ensembles:
                    Connection(
                        self.inhibit, e, transform=-self._inhib_scale, synapse=None
                    )
