# -*- coding: utf-8 -*-
#   This work is part of the Core Imaging Library developed by
#   Visual Analytics and Imaging System Group of the Science Technology
#   Facilities Council, STFC

#   Copyright 2018 Jakob Jorgensen, Daniil Kazantsev and Edoardo Pasca

#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at

#       http://www.apache.org/licenses/LICENSE-2.0

#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

from ccpi.optimisation.ops import Identity, FiniteDiff2D
import numpy
from ccpi.framework import DataContainer
import warnings
from ccpi.optimisation.functions import Function
def isSizeCorrect(data1 ,data2):
    if issubclass(type(data1), DataContainer) and \
       issubclass(type(data2), DataContainer):
        # check dimensionality
        if data1.check_dimensions(data2):
            return True
    elif issubclass(type(data1) , numpy.ndarray) and \
         issubclass(type(data2) , numpy.ndarray):
        return data1.shape == data2.shape
    else:
        raise ValueError("{0}: getting two incompatible types: {1} {2}"\
                         .format('Function', type(data1), type(data2)))
    return False
class Norm2(Function):
    
    def __init__(self, 
                 gamma=1.0, 
                 direction=None):
        super(Norm2, self).__init__()
        self.gamma     = gamma;
        self.direction = direction; 
    
    def __call__(self, x, out=None):
        
        if out is None:
            xx = numpy.sqrt(numpy.sum(numpy.square(x.as_array()), self.direction,
                                  keepdims=True))
        else:
            if isSizeCorrect(out, x):
                # check dimensionality
                if issubclass(type(out), DataContainer):
                    arr = out.as_array()
                    numpy.square(x.as_array(), out=arr)
                    xx = numpy.sqrt(numpy.sum(arr, self.direction, keepdims=True))
                        
                elif issubclass(type(out) , numpy.ndarray):
                    numpy.square(x.as_array(), out=out)
                    xx = numpy.sqrt(numpy.sum(out, self.direction, keepdims=True))
            else:
                raise ValueError ('Wrong size: x{0} out{1}'.format(x.shape,out.shape) )
        
        p  = numpy.sum(self.gamma*xx)        
        
        return p
    
    def prox(self, x, tau):

        xx = numpy.sqrt(numpy.sum( numpy.square(x.as_array()), self.direction, 
                                  keepdims=True ))
        xx = numpy.maximum(0, 1 - tau*self.gamma / xx)
        p  = x.as_array() * xx
        
        return type(x)(p,geometry=x.geometry)
    def proximal(self, x, tau, out=None):
        if out is None:
            return self.prox(x,tau)
        else:
            if isSizeCorrect(out, x):
                # check dimensionality
                if issubclass(type(out), DataContainer):
                    numpy.square(x.as_array(), out = out.as_array())
                    xx = numpy.sqrt(numpy.sum( out.as_array() , self.direction, 
                                  keepdims=True ))
                    xx = numpy.maximum(0, 1 - tau*self.gamma / xx)
                    x.multiply(xx, out= out.as_array())
                    
                        
                elif issubclass(type(out) , numpy.ndarray):
                    numpy.square(x.as_array(), out=out)
                    xx = numpy.sqrt(numpy.sum(out, self.direction, keepdims=True))
                    
                    xx = numpy.maximum(0, 1 - tau*self.gamma / xx)
                    x.multiply(xx, out= out)
            else:
                raise ValueError ('Wrong size: x{0} out{1}'.format(x.shape,out.shape) )
        

class TV2D(Norm2):
    
    def __init__(self, gamma):
        super(TV2D,self).__init__(gamma, 0)
        self.op = FiniteDiff2D()
        self.L = self.op.get_max_sing_val()
        

# Define a class for squared 2-norm
class Norm2sq(Function):
    '''
    f(x) = c*||A*x-b||_2^2
    
    which has 
    
    grad[f](x) = 2*c*A^T*(A*x-b)
    
    and Lipschitz constant
    
    L = 2*c*||A||_2^2 = 2*s1(A)^2
    
    where s1(A) is the largest singular value of A.
    
    '''
    
    def __init__(self,A,b,c=1.0,memopt=False):
        super(Norm2sq, self).__init__()
    
        self.A = A  # Should be an operator, default identity
        self.b = b  # Default zero DataSet?
        self.c = c  # Default 1.
        if memopt:
            try:
                self.range_tmp = A.range_geometry().allocate()
                self.domain_tmp = A.domain_geometry().allocate()
                self.memopt = True
            except NameError as ne:
                warnings.warn(str(ne))
                self.memopt = False
            except NotImplementedError as nie:
                print (nie)
                warnings.warn(str(nie))
                self.memopt = False
        else:
            self.memopt = False
        
        # Compute the Lipschitz parameter from the operator if possible
        # Leave it initialised to None otherwise
        try:
            self.L = 2.0*self.c*(self.A.norm()**2)
        except AttributeError as ae:
            pass
        except NotImplementedError as noe:
            pass
        
    #def grad(self,x):
    #    return self.gradient(x, out=None)

    def __call__(self,x):
        #return self.c* np.sum(np.square((self.A.direct(x) - self.b).ravel()))
        #if out is None:
        #    return self.c*( ( (self.A.direct(x)-self.b)**2).sum() )
        #else:
        y = self.A.direct(x)
        y.__isub__(self.b)
        #y.__imul__(y)
        #return y.sum() * self.c
        try:
            return y.squared_norm() * self.c
        except AttributeError as ae:
            # added for compatibility with SIRF 
            return (y.norm()**2) * self.c
    
    def gradient(self, x, out = None):
        if self.memopt:
            #return 2.0*self.c*self.A.adjoint( self.A.direct(x) - self.b )
            
            self.A.direct(x, out=self.range_tmp)
            self.range_tmp -= self.b 
            self.A.adjoint(self.range_tmp, out=out)
            #self.direct_placehold.multiply(2.0*self.c, out=out)
            out *= (self.c * 2.0)
        else:
            return (2.0*self.c)*self.A.adjoint( self.A.direct(x) - self.b )



# Box constraints indicator function. Calling returns 0 if argument is within 
# the box. The prox operator is projection onto the box. Only implements one 
# scalar lower and one upper as constraint on all elements. Should generalise 
# to vectors to allow different constraints one elements.
class IndicatorBox(Function):
    
    def __init__(self,lower=-numpy.inf,upper=numpy.inf):
        # Do nothing
        super(IndicatorBox, self).__init__()
        self.lower = lower
        self.upper = upper
        
    
    def __call__(self,x):
        
        if (numpy.all(x.array>=self.lower) and 
            numpy.all(x.array <= self.upper) ):
            val = 0
        else:
            val = numpy.inf
        return val
    
    def prox(self,x,tau=None):
        return  (x.maximum(self.lower)).minimum(self.upper)
    
    def proximal(self, x, tau, out=None):
        if out is None:
            return self.prox(x, tau)
        else:
            if not x.shape == out.shape:
                raise ValueError('Norm1 Incompatible size:',
                                 x.shape, out.shape)
            #(x.abs() - tau*self.gamma).maximum(0) * x.sign()
            x.abs(out = out)
            out.__isub__(tau*self.gamma)
            out.maximum(0, out=out)
            if self.sign_x is None or not x.shape == self.sign_x.shape:
                self.sign_x = x.sign()
            else:
                x.sign(out=self.sign_x)
                
            out.__imul__( self.sign_x )

# A more interesting example, least squares plus 1-norm minimization.
# Define class to represent 1-norm including prox function
class Norm1(Function):
    
    def __init__(self,gamma):
        super(Norm1, self).__init__()
        self.gamma = gamma
        self.L = 1
        self.sign_x = None
    
    def __call__(self,x,out=None):
        if out is None:
            return self.gamma*(x.abs().sum())
        else:
            if not x.shape == out.shape:
                raise ValueError('Norm1 Incompatible size:',
                                 x.shape, out.shape)
            x.abs(out=out)
            return out.sum() * self.gamma
    
    def prox(self,x,tau):
        return (x.abs() - tau*self.gamma).maximum(0) * x.sign()
    
    def proximal(self, x, tau, out=None):
        if out is None:
            return self.prox(x, tau)
        else:
            if isSizeCorrect(x,out):
                # check dimensionality
                if issubclass(type(out), DataContainer):
                    v = (x.abs() - tau*self.gamma).maximum(0)
                    x.sign(out=out)
                    out *= v
                    #out.fill(self.prox(x,tau))    
                elif issubclass(type(out) , numpy.ndarray):
                    v = (x.abs() - tau*self.gamma).maximum(0)
                    out[:] = x.sign()
                    out *= v
                    #out[:] = self.prox(x,tau)
            else:
                raise ValueError ('Wrong size: x{0} out{1}'.format(x.shape,out.shape) )
