from scipy.optimize import OptimizeResult as _OptimizeResult # For mocking SciPy OptimizeResult
import numpy as _numpy # Vectorized number crunching
from ..other_tools import progress as _progress # For monitoring progress
from .optimizer import Optimizer as _Optimizer # Optimizer parent class

"""
SciPy-style stochastic gradient descent algorithms

- Based on https://gist.github.com/jcmgray
- My future plan is to migrate this from pure Python code to a Nim wrapper
"""

""" Stochastic gradient descent functions """
# Standard stochastic gradient descent
def _sgd(fcn,
         x0,
         jac,
         args = (),
         learning_rate = 1e-3,
         maxiter = 1000,
         callback = None,
         verbose = False,
         learning_rate_schedule = None,
         nesterov = False,
         **kwargs
         ):
    """ 
    Stochastic gradient descent (SGD):

    - Bare bones implementation of SGD
    - Herbert R., Sutton M., "A Stochastic Approximation Method", JSTOR
      - https://www.jstor.org/stable/2236626
    """
    # Prepare learning rate schedule
    if learning_rate_schedule is None: learning_rate_schedule = lambda itn, lr: lr
    
    # Do stochastic gradient descent
    x = x0
    for iteration in range(0, maxiter):
        if callback and callback(iteration, x): break
        if verbose: _progress.update_progress((iteration + 1) / maxiter)
        else: x -= learning_rate_schedule(iteration, learning_rate) * jac(x, *args)
            
    # Return result of stochastic gradient descent
    return _OptimizeResult(
        x = x,
        fun = fcn(x, *args),
        jac = jac(x, *args),
        nit = iteration + 1,
        nfev = iteration + 1,
        success = True
    )

# Original version of "Adam" variant of stochastic gradient descent
def _adam(fcn,
          x0,
          jac,
          args = (),
          learning_rate = 1e-3,  # <---
          eps = 1e-8,            #       TensorFlow & PyTorch defaults
          beta1 = 0.9,           #      Also recommended in orig. paper
          beta2 = 0.999,         # <---
          maxiter = 1000,
          momentum_decay = 4e-3, # <--- Default in PyTorch
          callback = None,
          verbose = False,
          nesterov = True, # For "Nadam" - Adam w/ Nesterov momentum
          learning_rate_schedule = None, # User-specifified learning rate decay
          **kwargs
          ):
    """
    Adam variant of SGD:

    - Diederik P.K., Jimmy B., "Adam: A Method for Stochastic Optimization"
      - https://arxiv.org/abs/1412.6980
    """
    # Prepare variables for Adam optimization
    x = x0; j = _numpy.zeros(len(jac(x)));
    m = _numpy.zeros(len(x0)); v = _numpy.zeros(len(x0));
    mhat = _numpy.zeros(len(x0)); vhat = _numpy.zeros(len(x0));

    # Prepare some quantites to reduce repeated calculations
    beta1p = beta1; beta2p = beta2;
    mbeta1 = 1. - beta1; mbeta2 = 1. - beta2;
    mut = 0.; mutp1 = 0.;
    
    # Prepare learning rate schedule
    if learning_rate_schedule is None: learning_rate_schedule = lambda itn, lr: lr
    
    # Do Adam optimization
    for iteration in range(1, maxiter):
        if verbose: _progress.update_progress((iteration + 1) / maxiter)
        if callback and callback(iteration + 1, x): break
        j = jac(x, *args); m = beta1 * m + mbeta1 * j; v = beta2 * v + mbeta2 * j * j;
        if nesterov:
            mut = beta1 * (1. - 0.5 * _numpy.power(0.96, iteration * momentum_decay))
            mutp1 = beta1 * (1. - 0.5 * _numpy.power(0.96, (iteration + 1.) * momentum_decay))
            mhat = mutp1 * m / (1. - beta1p * beta1) + (1. - mut) * j / (1. - beta1p)
        else: mhat = m / (1. - beta1p)
        vhatroot = _numpy.sqrt( v / (1. - beta2p) ) + eps
        x -= learning_rate_schedule(iteration, learning_rate) * mhat / vhatroot
        beta1p *= beta1; beta2p *= beta2;
        
    # Return result of Adam
    return _OptimizeResult(
        x = x,
        fun = fcn(x, *args),
        jac = jac(x, *args),
        nit = iteration + 1,
        nfev = iteration + 1,
        success = True
    )
    
""" Stochastic gradient descent class """
class StochasticGradientDescent(_Optimizer):
    def __init__(self,
                 fcn = None,
                 jac = None,
                 optimizer = 'adam',
                 optimizer_arguments = {},
                 ):
        super().__init__(fcn = fcn, optimizer_arguments = optimizer_arguments, jac = jac)
        self._optimizer = optimizer
    def __call__(self, p0):
        match self._optimizer:
            case 'sgd': return _sgd(self._fcn, p0, self._jac, **self._args)
            case 'adam': return _adam(self._fcn, p0, self._jac, **self._args)
            case _: print(self._optimizer, 'is not a valid SGD optimizer choice.')
