
# Authors: Thierry Moudiki
#
# License: BSD 3 Clear


import numpy as np
from sklearn.base import BaseEstimator
from . import _optimizerc as optimizerc


class Optimizer(BaseEstimator):
    """Optimizer class
    
        Parameters
       ----------
       type_optim: str
           type of optimizer, (currently) either 'sgd' (stochastic minibatch gradient descent)
           or 'scd' (stochastic minibatch coordinate descent)
       num_iters: int
           number of iterations of the optimizer
       learning_rate: float
           step size
       batch_prop: float
           proportion of the initial data used at each optimization step
       learning_method: str
           "poly" - learning rate decreasing as a polynomial function 
           of # of iterations (default)
           "exp" - learning rate decreasing as an exponential function 
           of # of iterations
           "momentum" - gradient descent using momentum
       randomization: str
           type of randomization applied at each step
           "strat" - stratified subsampling (default)
           "shuffle" - random subsampling
       mass: float
           mass on velocity, for `method` == "momentum"
       decay: float
           coefficient of decrease of the learning rate for 
           `method` == "poly" and `method` == "exp"
       verbose: int 
           controls verbosity of gradient descent
           0 - nothing is printed
           1 - a progress bar is printed
           2 - succesive loss function values are printed
    """
    
    
    # construct the object -----

    def __init__(
        self,        
        type_optim="sgd",
        num_iters=100,         
        learning_rate=0.01, 
        batch_prop=1.0,         
        learning_method="momentum", 
        randomization="strat", 
        mass=0.9, decay=0.1, 
        verbose=1):
        
        self.type_optim = type_optim
        self.num_iters = num_iters
        self.learning_rate = learning_rate
        self.batch_prop = batch_prop
        self.learning_method = learning_method
        self.randomization = randomization
        self.mass = mass
        self.decay = decay
        self.verbose = verbose
        self.opt = None
                                                          
    
    def fit(self, loss_func, response, x0, **kwargs):
        """Fit GLM model to training data (X, y).
        
        Parameters
        ----------
        
        loss_func: loss function
        
        response: array-like, shape = [n_samples]
           target variable (used for subsampling) 
        
        x0: array-like, shape = [n_features]
            initial value provided to the optimizer
    
        **kwargs: additional parameters to be passed to 
                  loss function
               
        Returns
        -------
        self: object
        """
        
        
        if self.type_optim == "scd":
                    
            self.results = optimizerc.scd(loss_func, 
                            response=response, x=x0, 
                            num_iters=self.num_iters,
                            batch_prop=self.batch_prop, 
                            learning_rate=self.learning_rate, 
                            learning_method=self.learning_method, 
                            mass=self.mass, decay=self.decay,                     
                            randomization=self.randomization,             
                            verbose=self.verbose, 
                            **kwargs)            
            
            
        if self.type_optim == "sgd":
                                        
            self.results = optimizerc.sgd(loss_func, 
                            response=response, x=x0, 
                            num_iters=self.num_iters,
                            batch_prop=self.batch_prop, 
                            learning_rate=self.learning_rate, 
                            learning_method=self.learning_method, 
                            mass=self.mass, decay=self.decay,                     
                            randomization=self.randomization,             
                            verbose=self.verbose, 
                            **kwargs)    

        return self

    def one_hot_encode(self, y, n_classes):
        return optimizerc.one_hot_encode(y, n_classes)