from copy import copy
import warnings

import numpy as np

import nengo
from nengo.config import SupportDefaultsMixin
from nengo.exceptions import NotAddedToNetworkWarning, ValidationError
from nengo.params import (
    FrozenObject,
    IntParam,
    iter_params,
    NumberParam,
    Parameter,
    StringParam,
    Unconfigurable,
)
from nengo.utils.compat import is_integer, iteritems, range, with_metaclass
from nengo.utils.numpy import as_shape, maxint


class NetworkMember(type):
    """A metaclass used to add instances of derived classes to networks.

    Inheriting from this class means that Network.add will be invoked after
    initializing the object, unless add_to_container=False is passed to the
    derived class constructor.
    """

    def __call__(cls, *args, **kwargs):
        """Override default __call__ behavior so that Network.add is called."""
        inst = cls.__new__(cls)
        add_to_container = kwargs.pop('add_to_container', True)
        # Do the __init__ before adding in case __init__ errors out
        inst.__init__(*args, **kwargs)
        if add_to_container:
            nengo.Network.add(inst)
        inst._initialized = True  # value doesn't matter, just existence
        return inst


class NengoObject(with_metaclass(NetworkMember, SupportDefaultsMixin)):
    """A base class for Nengo objects.

    Parameters
    ----------
    label : string
        A descriptive label for the object.
    seed : int
        The seed used for random number generation.

    Attributes
    ----------
    label : string
        A descriptive label for the object.
    seed : int
        The seed used for random number generation.
    """

    # Order in which parameters have to be initialized.
    # Missing parameters will be initialized last in an undefined order.
    # This is needed for pickling and copying of Nengo objects when the
    # parameter initialization order matters.
    _param_init_order = []

    label = StringParam('label', default=None, optional=True)
    seed = IntParam('seed', default=None, optional=True)

    def __init__(self, label, seed):
        super(NengoObject, self).__init__()
        self.label = label
        self.seed = seed

    def __getstate__(self):
        state = self.__dict__.copy()
        del state['_initialized']

        for attr in self.params:
            param = getattr(type(self), attr)
            if self in param:
                state[attr] = getattr(self, attr)

        return state

    def __setstate__(self, state):
        for attr in self._param_init_order:
            setattr(self, attr, state.pop(attr))

        for attr in self.params:
            if attr in state:
                setattr(self, attr, state.pop(attr))

        for k, v in iteritems(state):
            setattr(self, k, v)

        self._initialized = True
        if len(nengo.Network.context) > 0:
            warnings.warn(NotAddedToNetworkWarning(self))

    def __setattr__(self, name, val):
        if hasattr(self, '_initialized') and not hasattr(self, name):
            warnings.warn(
                "Creating new attribute '%s' on '%s'. "
                "Did you mean to change an existing attribute?" % (name, self),
                SyntaxWarning)
        super(NengoObject, self).__setattr__(name, val)

    def __str__(self):
        return self._str(
            include_id=not hasattr(self, 'label') or self.label is None)

    def __repr__(self):
        return self._str(include_id=True)

    def _str(self, include_id):
        return "<%s%s%s>" % (
            type(self).__name__,
            "" if not hasattr(self, 'label') else
            " (unlabeled)" if self.label is None else
            ' "%s"' % self.label,
            " at 0x%x" % id(self) if include_id else "")

    @property
    def params(self):
        """Returns a list of parameter names that can be set."""
        return list(iter_params(self))

    def copy(self, add_to_container=True):
        with warnings.catch_warnings():
            # We warn when copying since we can't change add_to_container.
            # However, we deal with it here, so we ignore the warning.
            warnings.simplefilter('ignore', category=NotAddedToNetworkWarning)
            c = copy(self)
        if add_to_container:
            nengo.Network.add(c)
        return c


class ObjView(object):
    """Container for a slice with respect to some object.

    This is used by the __getitem__ of Neurons, Node, and Ensemble, in order
    to pass slices of those objects to Connection. This is a notational
    convenience for creating transforms. See Connection for details.

    Does not currently support any other view-like operations.
    """

    def __init__(self, obj, key=slice(None)):
        self.obj = obj

        # Node.size_in != size_out, so one of these can be invalid
        # NumPy <= 1.8 raises a ValueError instead of an IndexError.
        try:
            self.size_in = np.arange(self.obj.size_in)[key].size
        except (IndexError, ValueError):
            self.size_in = None
        try:
            self.size_out = np.arange(self.obj.size_out)[key].size
        except (IndexError, ValueError):
            self.size_out = None
        if self.size_in is None and self.size_out is None:
            raise IndexError("Invalid slice '%s' of %s" % (key, self.obj))

        if is_integer(key):
            # single slices of the form [i] should be cast into
            # slice objects for convenience
            if key == -1:
                # special case because slice(-1, 0) gives the empty list
                self.slice = slice(key, None)
            else:
                self.slice = slice(key, key+1)
        else:
            self.slice = key

    def copy(self):
        return copy(self)

    def __len__(self):
        return self.size_out

    def __str__(self):
        return "%s[%s]" % (self.obj, self._slice_string)

    def __repr__(self):
        return "%r[%s]" % (self.obj, self._slice_string)

    @property
    def _slice_string(self):
        if isinstance(self.slice, slice):
            sl_start = "" if self.slice.start is None else self.slice.start
            sl_stop = "" if self.slice.stop is None else self.slice.stop
            if self.slice.step is None:
                return "%s:%s" % (sl_start, sl_stop)
            else:
                return "%s:%s:%s" % (sl_start, sl_stop, self.slice.step)
        else:
            return str(self.slice)


class NengoObjectParam(Parameter):
    def __init__(self, name, optional=False, readonly=True,
                 nonzero_size_in=False, nonzero_size_out=False):
        default = Unconfigurable  # These can't have defaults
        self.nonzero_size_in = nonzero_size_in
        self.nonzero_size_out = nonzero_size_out
        super(NengoObjectParam, self).__init__(
            name, default, optional, readonly)

    def coerce(self, instance, nengo_obj):
        nengo_objects = (
            NengoObject,
            ObjView,
            nengo.ensemble.Neurons,
            nengo.connection.LearningRule
        )
        if not isinstance(nengo_obj, nengo_objects):
            raise ValidationError("'%s' is not a Nengo object" % nengo_obj,
                                  attr=self.name, obj=instance)
        if self.nonzero_size_in and nengo_obj.size_in < 1:
            raise ValidationError("'%s' must have size_in > 0." % nengo_obj,
                                  attr=self.name, obj=instance)
        if self.nonzero_size_out and nengo_obj.size_out < 1:
            raise ValidationError("'%s' must have size_out > 0." % nengo_obj,
                                  attr=self.name, obj=instance)
        return super(NengoObjectParam, self).coerce(instance, nengo_obj)


class Process(FrozenObject):
    """A general system with input, output, and state.

    For more details on how to use processes and make
    custom process subclasses, see :doc:`examples/advanced/processes`.

    Parameters
    ----------
    default_size_in : int (Default: 0)
        Sets the default size in for nodes using this process.
    default_size_out : int (Default: 1)
        Sets the default size out for nodes running this process. Also,
        if ``d`` is not specified in `~.Process.run` or `~.Process.run_steps`,
        this will be used.
    default_dt : float (Default: 0.001 (1 millisecond))
        If ``dt`` is not specified in `~.Process.run`, `~.Process.run_steps`,
        `~.Process.ntrange`, or `~.Process.trange`, this will be used.
    seed : int, optional (Default: None)
        Random number seed. Ensures random factors will be the same each run.

    Attributes
    ----------
    default_dt : float
        If ``dt`` is not specified in `~.Process.run`, `~.Process.run_steps`,
        `~.Process.ntrange`, or `~.Process.trange`, this will be used.
    default_size_in : int
        The default size in for nodes using this process.
    default_size_out : int
        The default size out for nodes running this process. Also, if ``d`` is
        not specified in `~.Process.run` or `~.Process.run_steps`,
        this will be used.
    seed : int or None
        Random number seed. Ensures random factors will be the same each run.
    """

    default_size_in = IntParam('default_size_in', low=0)
    default_size_out = IntParam('default_size_out', low=0)
    default_dt = NumberParam('default_dt', low=0, low_open=True)
    seed = IntParam('seed', low=0, high=maxint, optional=True)

    def __init__(self, default_size_in=0, default_size_out=1,
                 default_dt=0.001, seed=None):
        super(Process, self).__init__()
        self.default_size_in = default_size_in
        self.default_size_out = default_size_out
        self.default_dt = default_dt
        self.seed = seed

    def apply(self, x, d=None, dt=None, rng=np.random, copy=True, **kwargs):
        """Run process on a given input.

        Keyword arguments that do not appear in the parameter list below
        will be passed to the ``make_step`` function of this process.

        Parameters
        ----------
        x : ndarray
            The input signal given to the process.
        d : int, optional (Default: None)
            Output dimensionality. If None, ``default_size_out`` will be used.
        dt : float, optional (Default: None)
            Simulation timestep. If None, ``default_dt`` will be used.
        rng : `numpy.random.RandomState` (Default: ``numpy.random``)
            Random number generator used for stochstic processes.
        copy : bool, optional (Default: True)
            If True, a new output array will be created for output.
            If False, the input signal ``x`` will be overwritten.
        """
        shape_in = as_shape(np.asarray(x[0]).shape, min_dim=1)
        shape_out = as_shape(self.default_size_out if d is None else d)
        dt = self.default_dt if dt is None else dt
        rng = self.get_rng(rng)
        step = self.make_step(shape_in, shape_out, dt, rng, **kwargs)
        output = np.zeros((len(x),) + shape_out) if copy else x
        for i, xi in enumerate(x):
            output[i] = step((i+1) * dt, xi)
        return output

    def get_rng(self, rng):
        """Get a properly seeded independent RNG for the process step.

        Parameters
        ----------
        rng : `numpy.random.RandomState`
            The parent random number generator to use if the seed is not set.
        """
        seed = rng.randint(maxint) if self.seed is None else self.seed
        return np.random.RandomState(seed)

    def make_step(self, shape_in, shape_out, dt, rng):
        """Create function that advances the process forward one time step.

        This must be implemented by all custom processes.

        Parameters
        ----------
        shape_in : tuple
            The shape of the input signal.
        shape_out : tuple
            The shape of the output signal.
        dt : float
            The simulation timestep.
        rng : `numpy.random.RandomState`
            A random number generator.
        """
        raise NotImplementedError("Process must implement `make_step` method.")

    def run(self, t, d=None, dt=None, rng=np.random, **kwargs):
        """Run process without input for given length of time.

        Keyword arguments that do not appear in the parameter list below
        will be passed to the ``make_step`` function of this process.

        Parameters
        ----------
        t : float
            The length of time to run.
        d : int, optional (Default: None)
            Output dimensionality. If None, ``default_size_out`` will be used.
        dt : float, optional (Default: None)
            Simulation timestep. If None, ``default_dt`` will be used.
        rng : `numpy.random.RandomState` (Default: ``numpy.random``)
            Random number generator used for stochstic processes.
        """
        dt = self.default_dt if dt is None else dt
        n_steps = int(np.round(float(t) / dt))
        return self.run_steps(n_steps, d=d, dt=dt, rng=rng, **kwargs)

    def run_steps(self, n_steps, d=None, dt=None, rng=np.random, **kwargs):
        """Run process without input for given number of steps.

        Keyword arguments that do not appear in the parameter list below
        will be passed to the ``make_step`` function of this process.

        Parameters
        ----------
        n_steps : int
            The number of steps to run.
        d : int, optional (Default: None)
            Output dimensionality. If None, ``default_size_out`` will be used.
        dt : float, optional (Default: None)
            Simulation timestep. If None, ``default_dt`` will be used.
        rng : `numpy.random.RandomState` (Default: ``numpy.random``)
            Random number generator used for stochstic processes.
        """
        shape_in = as_shape(0)
        shape_out = as_shape(self.default_size_out if d is None else d)
        dt = self.default_dt if dt is None else dt
        rng = self.get_rng(rng)
        step = self.make_step(shape_in, shape_out, dt, rng, **kwargs)
        output = np.zeros((n_steps,) + shape_out)
        for i in range(n_steps):
            output[i] = step((i+1) * dt)
        return output

    def ntrange(self, n_steps, dt=None):
        """Create time points corresponding to a given number of steps.

        Parameters
        ----------
        n_steps : int
            The given number of steps.
        dt : float, optional (Default: None)
            Simulation timestep. If None, ``default_dt`` will be used.
        """
        dt = self.default_dt if dt is None else dt
        return dt * np.arange(1, n_steps + 1)

    def trange(self, t, dt=None):
        """Create time points corresponding to a given length of time.

        Parameters
        ----------
        t : float
            The given length of time.
        dt : float, optional (Default: None)
            Simulation timestep. If None, ``default_dt`` will be used.
        """
        dt = self.default_dt if dt is None else dt
        n_steps = int(np.round(float(t) / dt))
        return self.ntrange(n_steps, dt=dt)


class ProcessParam(Parameter):
    """Must be a Process."""

    def coerce(self, instance, process):
        self.check_type(instance, process, Process)
        return super(ProcessParam, self).coerce(instance, process)
