# Online stochastic gradient descent is a variant of 
# stochastic gradient descent in which you estimate the 
# gradient of the cost function for each observation and 
# update the decision variables accordingly. 
#
#  This can help you find the global minimum, especially if 
# the objective function is convex.

# Batch stochastic gradient descent is somewhere between 
# ordinary gradient descent and the online method. 
# The gradients are calculated and the decision variables 
# are updated iteratively with subsets of all observations, 
# called minibatches. 
#
# Epoch = one iteration over all the observations
#       = all mini-batches
#
#  This variant is very popular for training neural networks.

#########################################################################
# You can imagine the online algorithm as a special kind of 
# batch algorithm in which each minibatch has only one observation. 
#
# Classical gradient descent is another special case in which there’s 
# only one batch containing all observations.
#########################################################################


import numpy as np

#######################################################################
# Stochatic Gradient Descent (with mini-batches)
#######################################################################

def sgd(gradient, x, y, start, learn_rate=0.1, batch_size=1, n_iter=50,
        tolerance=1e-06, dtype="float64", random_state=None):

    # Checking if the gradient is callable
    if not callable(gradient):
        raise TypeError("'gradient' must be callable")


    # Setting up the data type for NumPy arrays
    dtype_ = np.dtype(dtype)

    # Converting x and y to NumPy arrays
    x, y = np.array(x, dtype=dtype_), np.array(y, dtype=dtype_)

    n_obs = x.shape[0]

    if n_obs != y.shape[0]:
        raise ValueError("'x' and 'y' lengths do not match")

    xy = np.c_[x.reshape(n_obs, -1), y.reshape(n_obs, 1)]

    print()
    print("x = ", x)
    print("y = ", y)
    print("xy = ")
    print(xy)

    # Initializing the random number generator
    seed = None if random_state is None else int(random_state)

    rng = np.random.default_rng(seed=seed)


    # Initializing the values of the variables
    vector = np.array(start, dtype=dtype_)

    # Setting up and checking the learning rate
    learn_rate = np.array(learn_rate, dtype=dtype_)

    if np.any(learn_rate <= 0):
        raise ValueError("'learn_rate' must be greater than zero")


    # Setting up and checking the size of minibatches
    batch_size = int(batch_size)

    if not 0 < batch_size <= n_obs:
        raise ValueError(
            "'batch_size' must be greater than zero and less than "
            "or equal to the number of observations"
        )

    # Setting up and checking the maximal number of iterations
    n_iter = int(n_iter)

    if n_iter <= 0:
        raise ValueError("'n_iter' must be greater than zero")

    # Setting up and checking the tolerance
    tolerance = np.array(tolerance, dtype=dtype_)

    if np.any(tolerance <= 0):
        raise ValueError("'tolerance' must be greater than zero")

    ###############################################
    # Performing the gradient descent loop
    ###############################################
    for _ in range(n_iter):

        print(">>> Start of a EPOCH")
        # Shuffle x and y
        rng.shuffle(xy)

        print(">>> Random ordering:")
        print(xy)

        # Performing minibatch moves
        for start in range(0, n_obs, batch_size):

            stop = start + batch_size
            x_batch, y_batch = xy[start:stop, :-1], xy[start:stop, -1:]
            print(">>> Next batch:")
            print(x_batch, y_batch)

            # Recalculating the difference
            grad = np.array(gradient(x_batch, y_batch, vector), dtype_)
            print(">>> grad = ", grad)
            diff = -learn_rate * grad

            # Checking if the absolute difference is small enough
            if np.all(np.abs(diff) <= tolerance):
                break

            # Updating the values of the variables
            vector += diff

    print("++ vector.shape = ", vector.shape)
    return vector if vector.shape else vector.item()



 ###################################################################
# Gradient descent to find mim Sum of Squred Residuals (error function)
#
# ssr_gradient(x, y, b):
#
#    x and y, which contain the observation inputs and outputs (np.array()
#    b that holds the current values of the decision variables 𝑏₀ and 𝑏₁.
#
# Return: [best b[0] = mean(b[0] + b[1] * x[k] - y[k]),
#          best b[1] = mean( (b[0] + b[1] * x[k] - y[k]) x[k] )]

def ssr_gradient(x, y, b):

    print("G>> x,y,b = ", x, y, b)
    res = b[0] + b[1] * x - y           # Vector !!!
    z = res.mean(), (res * x).mean()    # [a, b]
    print("Gradient = ", z)
    return z
#   return res.mean(), (res * x).mean() # .mean() is a method of np.ndarray


x = np.array([5, 15, 25, 35, 45, 55])
y = np.array([5, 20, 14, 32, 22, 38])

out = sgd(ssr_gradient, x, y, start=[0.5, 0.5], learn_rate=0.0001,
          batch_size=3, n_iter=100_000, random_state=0)
print(out)
# array([5.63014443, 0.53901017])