"""
Constraints for Keras trainable weights
=======================================
"""

import scipy as sp
from keras import backend as K
from keras.constraints import Constraint
from theano import tensor as T
import logging
clip_logger = logging.getLogger("Clip")


def _my_clip(x, min_value, max_value):
    max_value = T.maximum(min_value, max_value)
    return T.clip(x, min_value, max_value)


# monkey patching
K.clip = _my_clip


class Clip(Constraint):
    """
    Clip weights within a range.

    Parameters
    ----------

    minimum: 2d array
       Minimum weights.

    maximum: 2d array
       Maximum weights.
    """
    def __init__(self, minimum, maximum):
        self.minimum = minimum
        self.maximum = maximum

    @classmethod
    def random_sparse(cls, shape, p, max_min=100):
        """
        Create random sparsity pattern.
        """

        non_zero = (sp.rand(*shape) < p).astype(float)
        return cls.from_unsigned_matrix(non_zero, max_min=max_min)

    @classmethod
    def from_signed_matrix(cls, matrix, max_min=100):
        """
        Preserve sparsity pattern and signs.

        Parameters
        ----------
        max_min: float
            Max and min of the allowed range.

        matrix: ndarray
            Positive entries are restricted to be positive.
            Negative entries are restricted to be negative.
            Zero entries are restricted to be zero.
        """
        clip_logger.debug("Build clip from signed matrix")
        maximum = sp.zeros_like(matrix)
        minimum = sp.zeros_like(matrix)
        maximum[matrix > 0] = max_min
        minimum[matrix < 0] = -max_min
        return cls(minimum, maximum)

    @classmethod
    def from_unsigned_matrix(cls, matrix, max_min=100):
        """
        Preserve sparsity pattern.

        Parameters
        ----------
        max_min: float
            Maximum and minimum of the allowed range.

        matrix: ndarray
            Zero entries are restricted to be zero.
        """
        maximum = sp.zeros_like(matrix)
        minimum = sp.zeros_like(matrix)
        maximum[matrix != 0] = max_min
        minimum[matrix != 0] = -max_min
        return cls(minimum, maximum)

    def __repr__(self):
        return "{}({})".format(self.__class__.__name__, self.minimum, self.maximum)

    def __call__(self, p):
        clip_logger.debug("Clip called")
        return K.clip(p, K.variable(self.minimum), K.variable(self.maximum))

    def get_config(self):
        return {"name": self.__class__.__name__,
                "minimum": self.minimum,
                "maximum": self.maximum}


