import numpy as np 
import autodiff.function as F
from autodiff.variable import Variable
from autodiff.variable import ReverseVariable
import matplotlib.pyplot as plt

"""
Optimizer class and subclasses that uses automatic differentiations
to find local critical points.
Methods implemented: RMSProp, GradientDescent.
"""

class Optimizer:
    """
    Base class for optimizing methods:
        Implements base methods for a family of optimizers.
        _step method not implemented and base class can not optimize a given function.
    """
    def __init__(self, lr, tol, loss_fn, init_point):
        """
        init_point: numpy.array or autodiff.Variable
        loss_fn should be a function wrapped around autodiff module
            loss_fn input: N-Dimensional Variable
            loss_fn output: 1-Dimensional Variable.
            Should be positive by definition of a loss function.
            Any type of function could be used.
        Remark: We enforce the loss_fn to be a 1-Dimensional,
        because the space of real numbers is a implicit ordered set.
        lr: strictly positive float
        tol: strictly positive float-Stops the iterations whenever the loss_fn
        becomes strictly inferior to tol.
        """
        assert isinstance(lr, (int, float)) and not isinstance(lr, bool), "lr should be numeric type"
        assert isinstance(tol, (int, float)) and not isinstance(
            tol, bool), "tol should be numeric type"

        assert lr > 0., "Need a positive learning rate"
        assert tol > 0., "Need a positive learning rate"
        self.lr = lr
        self.tol = tol
        self.loss_fn = loss_fn
        #Check whether we will be working with variable indeed.
        try:
            self.current_point = Variable(init_point)
        except TypeError as e:
            if isinstance(init_point, Variable):
                self.current_point = init_point
            else:
                raise TypeError(e)
        #Check whether the loss_fn is in the right space
        out = loss_fn(self.current_point)
        assert isinstance(
            out.val, float), "The loss function should be scalar-output"
    
    def _step(self, *args, **kwargs):
        """
        Given a loss Variable (with uni-dimensional value), 
        should update self.current_point (same dimensions for the value and gradient. ) 
        """
        raise NotImplementedError

    def _eval(self, *args, **kwargs):
        """
        Output is a Variable.
        Keep args if ever we want variable length inputs
        """
        return self.loss_fn(*args, **kwargs)
    
    def minimize(self, nb_steps, keep_track=True):
        """
        Keep track is a bool-> True returns the different losses/points obtained durring optim.
        """
        trajectory = []
        losses = []
        it = 0 
        loss = Variable(val=self.tol+1) #Randomly initialize the loss to get into the 
        while it < nb_steps and loss.val > self.tol: #Keeping the loss positive 
            loss = self._eval(self.current_point)
            #self.current_point -= self.lr * loss.grad #By doing this, we directly create a new Variable
            self._step(loss)
            #keep track of our thing
            trajectory.append(self.current_point.val)
            losses.append(loss.val)
            it +=1
        print('Minimized the function for {} steps.'.format(it))
        if keep_track:
            assert self.current_point.val.shape
            return self.current_point, losses, trajectory
        else:
            return self.current_point

    def __repr__(self):
        return str(vars(self))

    def visualize(self, losses, trajectory):
        """
        Assumes matplotlib.pyplot is imported
        losses and trajectory are supposed to be generated by minimize
        """
        assert len(losses)>0 and len(trajectory)>0, "Can not provide empty lists"
        nb_coordinates = trajectory[0].shape[0] #p,1->gets 1.
    
        fig, ax = plt.subplots(1,nb_coordinates+1, figsize=(20,5))
        ax[0].plot(losses)
        ax[0].set_xlabel('Iterations')
        ax[0].set_ylabel('Function values')
        ax[0].set_title('Decrease of the function with respect to iterations.')
        for i in range(1, nb_coordinates+1):
            trajectory_ = [traj[i-1] for traj in trajectory]
            ax[i].plot(trajectory_)
            ax[i].set_xlabel('Iterations')
            ax[i].set_title(
                'Coordinate {}'.format(i-1))
        plt.show()

class GradientDescent(Optimizer):
    def _step(self, loss):
        """
        Assumes loss has a grad attribute.
        """
        self.current_point -= self.lr * loss.grad.T
        
class RMSProp(Optimizer):
    """
    Implements RMSProp. Inherits from Optimizer.
    Overloads __init__ and _step methods
    """
    def __init__(self, *args, beta=0.9):
        super().__init__(*args)
        self.beta = beta

    def _step(self, loss, eps=10e-8):
        try:
            self.avg = self.beta * self.avg + (1 - self.beta) * loss.grad ** 2 #Loss val and grad should be (N,1) and (N,1)
        except Exception as e:#self.avg does not exist yet. Needs to create it. 
            #print(e)
            self.avg = np.zeros(loss.grad.shape, dtype=np.float64)
            self.avg = self.beta * self.avg + (1 - self.beta) * loss.grad ** 2
        #Update rule
        self.current_point -= self.lr * loss.grad.T / (np.sqrt(self.avg.T) + eps)#Element wise sqrt. Add eps for numerical overflow. 
  
#class Adam(Optimizer):
#   #TODO: or not.
#    def __init__(self, *args, beta1=0.9, beta2=0.99):
#        super().__init__(*args)
#        self.beta1 = beta1
#        self.beta2 = beta2

#    def _step(self, loss):
#        return NotImplementedError

