# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/000_common_core.ipynb.

# %% auto 0
__all__ = ['eval_accuracy', 'to_scalar', 'gg_confusion_matrix', 'Optimizer', 'SGD', 'Scheduler', 'ConstantLR',
           'CosineAnnealingLR', 'DataLoader', 'gradient_checker', 'J_theta_global', 'dJ_theta_global']

# %% ../nbs/000_common_core.ipynb 2
from .core import *
import numpy as np
from sklearn.metrics import accuracy_score
from sklearn.metrics import confusion_matrix
import pandas as pd
from lets_plot import *
from typing import List, Optional, Union, Callable


def eval_accuracy(model, val, y_val):
    model.eval()
    output = model(Tensor(val))
    y_hat = np.argmax(output.data, axis=1)
    model.train()
    return accuracy_score(y_val.ravel(), y_hat.ravel())


def to_scalar(x, i):
    """Returns i-th element from an array x if x is array, else returns x
Args:
    x - input array
    i - index to return
Returns:
    x[i] if x is an array and x else
"""
    if isinstance(x, (np.ndarray, list)):
        return x[i]
    return x

def gg_confusion_matrix(y, y_hat):
    conf_mat = confusion_matrix(y, y_hat)[::-1]
    confusion_dat = pd.DataFrame(conf_mat)
    observed = confusion_dat.columns.values
    actual = confusion_dat.index.values
    xx, yy = np.meshgrid(actual, observed)
    xx = xx.reshape(-1)
    yy = yy.reshape(-1)
    zz = conf_mat.reshape(-1)
    dat = {'predicted':xx, 'actual':yy[::-1], 'z':zz}
    p = ggplot(dat, aes('predicted', 'actual', fill='z')) \
        + geom_raster() \
        + geom_text(aes(label='z'), color='white')\
        + theme(legend_position='none', axis_ticks='blank', axis_line='blank')\
        + ggsize(500, 500) + scale_x_discrete() + scale_y_discrete()\
        + ggtitle('Confusion matrix')
    return p

# %% ../nbs/000_common_core.ipynb 4
class Optimizer:
    def __init__(self, params: List[Tensor], lr: float = 0.001):
        self.params = params
        self.lr = lr

    def step(self, *args, **kwargs):
        raise NotImplemented

    def zero_grad(self):
        for param in self.params:
            param.zero_grad()

# %% ../nbs/000_common_core.ipynb 5
class SGD(Optimizer):
    def __init__(self, params: List[Tensor], lr: float = 0.001):
        super().__init__(params, lr)

    def step(self):
        for param in self.params:
            param.data -= self.lr * param.grad

# %% ../nbs/000_common_core.ipynb 6
class Scheduler:
    def __init__(self, optimizer: Optimizer, last_epoch: int = -1):
        self.optimizer = optimizer
        self.base_lr = optimizer.lr
        self.last_epoch = last_epoch

    def step(self):
        raise NotImplementedError

# %% ../nbs/000_common_core.ipynb 7
class ConstantLR(Scheduler):
    def __init__(self, optimizer: Optimizer):
        super().__init__(optimizer)
        self.lr = optimizer.lr

    def step(self):
        self.last_epoch += 1

# %% ../nbs/000_common_core.ipynb 8
class CosineAnnealingLR(Scheduler):
    def __init__(self, optimizer: Optimizer, T_max: int, eta_min: float = 0, anneal_epochs: int = None, last_epoch: int = -1):
        super().__init__(optimizer)
        self.T_max = T_max
        self.eta_min = eta_min
        self.lr = optimizer.lr
        self.last_epoch = last_epoch
        self.start_epoch = last_epoch
        self.anneal_epochs = anneal_epochs

    @staticmethod
    def _cosine_anneal(t):
        return (1 + np.cos(np.pi * t)) / 2

    def get_lr(self):
        if self.anneal_epochs is not None:
            passed_epochs = self.last_epoch - self.start_epoch
            if passed_epochs > self.anneal_epochs:
                return self.lr
        t = self.last_epoch / self.T_max
        return self.eta_min + (self.base_lr - self.eta_min) * self._cosine_anneal(t)

    def step(self):
        self.lr = self.get_lr()
        self.optimizer.lr = self.lr
        self.last_epoch += 1

# %% ../nbs/000_common_core.ipynb 9
class DataLoader:
    def __init__(self, data, target, batch_size=20):
        self.data = data
        self.target = target
        self.batch_size = batch_size

    def next(self):
        m, _ = self.data.shape
        rand_index = np.random.choice(m, size=m, replace=False)
        X, y = self.data[rand_index], self.target[rand_index]
        pos = 0
        while pos < m:
            X_batch, y_batch = X[pos:pos+self.batch_size], y[pos:pos+self.batch_size]
            yield Tensor(X_batch, name="x"), Tensor(y_batch, name="y")
            pos += self.batch_size

    def __call__(self):
        return self.next()

# %% ../nbs/000_common_core.ipynb 11
def gradient_checker(J, grad_J, theta, eps=1e-5, rtol=1e-5):
    """Gradient checker for scalar and vector functions
Args:
    J - function of theta
    grad_J - gradient of function J
    theta - the point for which to compute the numerical gradient
    eps - step value in numerical gradient
    rtol - relative tolerance threshold value
Returns:
    error message if the relative tolerance is greater for some axis
    or "Gradient check passed" else
"""
    it = np.nditer(theta, flags=['multi_index'], op_flags=['readwrite'])

    while not it.finished:
        ix = it.multi_index
        theta_ = np.array(theta, copy=True)
        theta_[ix] += eps
        J1 = J(theta_)
        theta_ = np.array(theta, copy=True)
        theta_[ix] -= eps
        J2 = J(theta_)

        J1 = to_scalar(J1, ix)
        J2 = to_scalar(J2, ix)

        num_grad = (J1 - J2)/(2*eps)

        rel_tol = np.abs(num_grad - grad_J)[ix]/(1. + np.minimum(np.abs(num_grad), np.abs(grad_J[ix])))

        if rel_tol > rtol:
            print(f'num_grad: {num_grad} grad: {grad_J[ix]} factor: {grad_J[ix] / num_grad}')
            print(f'Incorrect gradient for the axis {str(ix)}')
            return
        it.iternext()
    print(f'Gradient check passed')

def J_theta_global(model, loss_function, theta, idx, x, y):
    original = model.parameters[idx].data.copy()
    np.copyto(dst=model.parameters[idx].data, src=theta)
    outputs = model(x)
    loss_value = loss_function(outputs, y).data
    np.copyto(dst=model.parameters[idx].data, src=original)
    model.zero_grad()
    return loss_value

def dJ_theta_global(model, loss_function, x, y):
    outputs = model(x)
    loss = loss_function(outputs, y)
    loss.backward()
    grads = []
    for parameter in model.parameters:
        grads.append(parameter.grad.copy())
    model.zero_grad()
    return grads
