import numpy as np

from .float import _f, _0, _1
from .optimizer_base import *
from .utility import *

# some of the the following optimizers are blatantly lifted from tiny-dnn:
# https://github.com/tiny-dnn/tiny-dnn/blob/master/tiny_dnn/optimizers/optimizer.h


class Momentum(Optimizer):
    def __init__(self, lr=0.01, mu=0.9, nesterov=False):
        self.mu = _f(mu)  # momentum
        self.nesterov = bool(nesterov)

        super().__init__(lr)

    def reset(self):
        self.Vprev = None

    def compute(self, dW, W):
        if self.Vprev is None:
            self.Vprev = np.copy(dW)

        V = self.mu * self.Vprev - self.lr * dW
        self.Vprev[:] = V
        if self.nesterov:
            return self.mu * V - self.lr * dW

        return V


class Adagrad(Optimizer):
    def __init__(self, lr=0.01, eps=1e-8):
        self.eps = _f(eps)

        super().__init__(lr)

    def reset(self):
        self.g = None

    def compute(self, dW, W):
        if self.g is None:
            self.g = np.zeros_like(dW)

        self.g += np.square(dW)
        return -self.lr * dW / (np.sqrt(self.g) + self.eps)


class RMSprop(Optimizer):
    # RMSprop generalizes* Adagrad, etc.

    # * RMSprop == Adagrad when
    #   RMSprop.mu == 1

    def __init__(self, lr=1e-4, mu=0.99, eps=1e-8):
        self.mu = _f(mu)  # decay term
        self.eps = _f(eps)

        # one might consider the following equation when specifying mu:
        # mu = e**(-1/t)
        # default: t = -1/ln(0.99) = ~99.5
        # therefore the default of mu=0.99 means
        # an input decays to 1/e its original amplitude over 99.5 batches.
        # (this is from DSP, so how relevant it is in SGD is debatable)

        super().__init__(lr)

    def reset(self):
        self.g = None

    def compute(self, dW, W):
        if self.g is None:
            self.g = np.zeros_like(dW)

        # basically apply a first-order low-pass filter to delta squared,
        self.g += (1 - self.mu) * (np.square(dW) - self.g)

        # and sqrt it to complete the running root-mean-square approximation.
        return -self.lr * dW / (np.sqrt(self.g) + self.eps)


class RMSpropCentered(Optimizer):
    # referenced TensorFlow/PyTorch.
    # paper: https://arxiv.org/pdf/1308.0850v5.pdf

    def __init__(self, lr=1e-4, aleph=0.95, momentum=0.9, eps=1e-8):
        self.aleph = _f(aleph)
        self.momentum = _f(momentum)
        self.eps = _f(eps)

        super().__init__(lr)

    def reset(self):
        self.g = None
        self.mt = None
        self.vt = None
        self.delta = None

    def compute(self, dW, W):
        if self.g is None:
            self.g = np.zeros_like(dW)
        if self.mt is None:
            self.mt = np.zeros_like(dW)
        if self.vt is None:
            self.vt = np.zeros_like(dW)
        if self.delta is None:
            self.delta = np.zeros_like(dW)

        self.mt += (1 - self.aleph) * (dW - self.mt)
        self.vt += (1 - self.aleph) * (np.square(dW) - self.vt)

        # PyTorch has the epsilon outside of the sqrt,
        # TensorFlow and the paper have it within.
        # in onn, we generally do it outside, as this seems to work better.
        temp = dW / (np.sqrt(self.vt - np.square(self.mt)) + self.eps)

        # TensorFlow does it this way.
        self.delta[:] = self.momentum * self.delta + self.lr * temp
        return -self.delta
        # PyTorch does it this way.
        # self.delta[:] = self.momentum * self.delta + temp
        # return -self.lr * self.delta
        # they are equivalent only when LR is constant, which it might not be.


class Adam(Optimizer):
    # paper: https://arxiv.org/abs/1412.6980
    # Adam generalizes* RMSprop, and
    # adds a decay term to the regular (non-squared) delta, and performs
    # debiasing to compensate for the filtered deltas starting from zero.

    # * Adam == RMSprop when
    #   Adam.b1 == 0
    #   Adam.b2 == RMSprop.mu

    def __init__(self, lr=0.002, b1=0.9, b2=0.999, eps=1e-8):
        self.b1 = _f(b1)  # decay term
        self.b2 = _f(b2)  # decay term
        self.b1_t_default = _f(b1)  # decay term power t
        self.b2_t_default = _f(b2)  # decay term power t
        self.eps = _f(eps)

        super().__init__(lr)

    def reset(self):
        self.mt = None
        self.vt = None
        self.b1_t = self.b1_t_default
        self.b2_t = self.b2_t_default

    def compute(self, dW, W):
        if self.mt is None:
            self.mt = np.zeros_like(dW)
        if self.vt is None:
            self.vt = np.zeros_like(dW)

        # decay gain
        self.b1_t *= self.b1
        self.b2_t *= self.b2

        # filter
        self.mt += (1 - self.b1) * (dW - self.mt)
        self.vt += (1 - self.b2) * (np.square(dW) - self.vt)

        return -self.lr * (self.mt / (1 - self.b1_t)) \
            / (np.sqrt(self.vt / (1 - self.b2_t)) + self.eps)


class Nadam(Optimizer):
    # paper: https://arxiv.org/abs/1412.6980
    # paper: http://cs229.stanford.edu/proj2015/054_report.pdf
    # TODO: double-check this implementation. also read the damn paper.
    # lifted from:
    # https://github.com/fchollet/keras/blob/5d38b04/keras/optimizers.py#L530
    # https://github.com/jpilaul/IFT6266_project/blob/master/Models/Algo_Momentum.py

    def __init__(self, lr=0.002, b1=0.9, b2=0.999, eps=1e-8):
        self.b1 = _f(b1)  # decay term
        self.b2 = _f(b2)  # decay term
        self.eps = _f(eps)

        super().__init__(lr)

    def reset(self):
        self.mt = None
        self.vt = None
        self.t = 0
        self.sched = 1

    def compute(self, dW, W):
        self.t += 1

        if self.mt is None:
            self.mt = np.zeros_like(dW)
        if self.vt is None:
            self.vt = np.zeros_like(dW)

        ut0 = self.b1 * (1 - 0.5 * 0.96**(self.t + 0))
        ut1 = self.b1 * (1 - 0.5 * 0.96**(self.t + 1))

        sched0 = self.sched * ut0
        sched1 = self.sched * ut0 * ut1
        self.sched = sched0

        gp = dW / (1 - sched0)

        self.mt += (1 - self.b1) * (dW - self.mt)
        self.vt += (1 - self.b2) * (np.square(dW) - self.vt)

        mtp = self.mt / (1 - sched1)
        vtp = self.vt / (1 - self.b2**self.t)

        mt_bar = (1 - ut0) * gp + ut1 * mtp

        return -self.lr * mt_bar / (np.sqrt(vtp) + self.eps)


# more

class FTML(Optimizer):
    # paper: http://www.cse.ust.hk/~szhengac/papers/icml17.pdf
    # author's implementation: https://github.com/szhengac/optim/commit/923555e

    def __init__(self, lr=0.0025, b1=0.6, b2=0.999, eps=1e-8):
        self.iterations = _0
        self.b1 = _f(b1)  # decay term
        self.b2 = _f(b2)  # decay term
        self.eps = _f(eps)

        super().__init__(lr)

    def reset(self):
        self.dt1 = None
        self.dt = None
        self.vt = None
        self.zt = None
        self.b1_t = _1
        self.b2_t = _1

    def compute(self, dW, W):
        if self.dt1 is None:
            self.dt1 = np.zeros_like(dW)
        if self.dt is None:
            self.dt = np.zeros_like(dW)
        if self.vt is None:
            self.vt = np.zeros_like(dW)
        if self.zt is None:
            self.zt = np.zeros_like(dW)

        # NOTE: we could probably rewrite these equations to avoid this copy.
        self.dt1[:] = self.dt[:]

        self.b1_t *= self.b1
        self.b2_t *= self.b2

        # hardly an elegant solution.
        lr = max(self.lr, self.eps)

        # same as Adam's vt.
        self.vt[:] = self.b2 * self.vt + (1 - self.b2) * dW * dW

        # you can factor "inner" out of Adam as well.
        inner = np.sqrt(self.vt / (1 - self.b2_t)) + self.eps
        self.dt[:] = (1 - self.b1_t) / lr * inner

        sigma_t = self.dt - self.b1 * self.dt1

        # Adam's mt minus the sigma term.
        self.zt[:] = self.b1 * self.zt + (1 - self.b1) * dW - sigma_t * W

        # subtract by weights to avoid having to override self.update.
        return -self.zt / self.dt - W


class MomentumClip(Optimizer):
    def __init__(self, lr=0.01, mu=0.9, nesterov=False, clip=1.0, debug=False):
        self.mu = _f(mu)
        self.clip = _f(clip)
        self.nesterov = bool(nesterov)
        self.debug = bool(debug)

        super().__init__(lr)

    def reset(self):
        self.accum = None

    def compute(self, dW, W):
        if self.accum is None:
            self.accum = np.zeros_like(dW)

        total_norm = np.linalg.norm(dW)
        clip_scale = self.clip / (total_norm + 1e-6)
        if clip_scale < 1:
            if self.debug:
                lament("clipping gradients; norm: {:10.5f}".format(total_norm))
            dW *= clip_scale

        self.accum[:] = self.accum * self.mu + dW
        if self.nesterov:
            return -self.lr * (self.accum * self.mu + dW)
        else:
            return -self.lr * self.accum


class YellowFin(Optimizer):
    # paper: https://arxiv.org/abs/1706.03471
    # knowyourmeme: http://cs.stanford.edu/~zjian/project/YellowFin/
    # author's implementation:
    # https://github.com/JianGoForIt/YellowFin/blob/master/tuner_utils/yellowfin.py
    # code lifted:
    # https://gist.github.com/botev/f8b32c00eafee222e47393f7f0747666

    def __init__(self, lr=0.1, mu=0.0, beta=0.999, window_size=20,
                 debias=True, clip=1.0):
        self.lr_default = _f(lr)
        self.mu_default = _f(mu)
        self.beta = _f(beta)
        self.window_size = int(window_size)  # curv_win_width
        self.debias_enabled = bool(debias)
        self.clip = _f(clip)

        self.mu = _f(mu)  # momentum
        super().__init__(lr)

    def reset(self):
        self.accum = None

        self.lr = self.lr_default
        self.mu = self.mu_default

        self.step = 0
        self.beta_t = self.beta

        self.curv_win = np.zeros([self.window_size, ], dtype=np.float32)

        self.h_min = None
        self.h_max = None

        self.g_lpf = 0
        # self.g_squared_lpf = 0
        self.g_norm_squared_lpf = 0
        self.g_norm_lpf = 0
        self.h_min_lpf = 0
        self.h_max_lpf = 0
        self.dist_lpf = 0
        self.lr_lpf = 0
        self.mu_lpf = 0

    def get_lr_mu(self):
        p = (np.square(self.dist_avg) * np.square(self.h_min)) \
            / (2 * self.g_var)
        w3 = p * (np.sqrt(0.25 + p / 27.0) - 0.5)
        w = np.power(w3, 1/3)
        y = w - p / (3 * w)
        sqrt_mu1 = y + 1

        sqrt_h_min = np.sqrt(self.h_min)
        sqrt_h_max = np.sqrt(self.h_max)
        sqrt_mu2 = (sqrt_h_max - sqrt_h_min) / (sqrt_h_max + sqrt_h_min)

        sqrt_mu = max(sqrt_mu1, sqrt_mu2)
        if sqrt_mu2 > sqrt_mu1:
            print('note: taking dr calculation. something may have exploded.')

        lr = np.square(1 - sqrt_mu) / self.h_min
        mu = np.square(sqrt_mu)
        return lr, mu

    def compute(self, dW, W):
        if self.accum is None:
            self.accum = np.zeros_like(dW)

        # TODO: prevent allocations everywhere by using [:].
        #       assuming that really works. i haven't actually checked.

        total_norm = np.linalg.norm(dW)
        clip_scale = self.clip / (total_norm + 1e-6)
        if clip_scale < 1:
            # print("clipping gradients; norm: {:10.5f}".format(total_norm))
            dW *= clip_scale

        # fmt = 'W std: {:10.7f}e-3,  dWstd: {:10.7f}e-3,  V std: {:10.7f}e-3'
        # print(fmt.format(np.std(W), np.std(dW) * 100, np.std(V) * 100))

        b = self.beta
        m1b = 1 - self.beta
        debias = 1 / (1 - self.beta_t) if self.debias_enabled else 1

        g = dW
        g_squared = np.square(g)
        g_norm_squared = np.sum(g_squared)
        g_norm = np.sqrt(g_norm_squared)

        self.curv_win[self.step % self.window_size] = g_norm_squared
        valid_window = self.curv_win[:min(self.window_size, self.step + 1)]
        h_min_t = np.min(valid_window)
        h_max_t = np.max(valid_window)

        self.g_lpf = b * self.g_lpf + m1b * g
        # self.g_squared_lpf = b * self.g_squared_lpf + m1b * g_squared
        self.g_norm_squared_lpf = b * self.g_norm_squared_lpf \
            + m1b * g_norm_squared
        self.g_norm_lpf = b * self.g_norm_lpf + m1b * g_norm
        self.h_min_lpf = b * self.h_min_lpf + m1b * h_min_t
        self.h_max_lpf = b * self.h_max_lpf + m1b * h_max_t

        g_avg = debias * self.g_lpf
        # g_squared_avg = debias * self.g_squared_lpf
        g_norm_squared_avg = debias * self.g_norm_squared_lpf
        g_norm_avg = debias * self.g_norm_lpf
        self.h_min = debias * self.h_min_lpf
        self.h_max = debias * self.h_max_lpf
        assert self.h_max >= self.h_min

        dist = g_norm_avg / g_norm_squared_avg

        self.dist_lpf = b * self.dist_lpf + m1b * dist

        self.dist_avg = debias * self.dist_lpf

        self.g_var = g_norm_squared_avg - np.sum(np.square(g_avg))
        # equivalently:
        # self.g_var = np.sum(np.abs(g_squared_avg - np.square(g_avg)))

        if self.step > 0:
            lr_for_real, mu_for_real = self.get_lr_mu()
            self.mu_lpf = b * self.mu_lpf + m1b * mu_for_real
            self.lr_lpf = b * self.lr_lpf + m1b * lr_for_real
            self.mu = debias * self.mu_lpf
            self.lr = debias * self.lr_lpf

        self.accum[:] = self.accum * self.mu - self.lr * dW
        V = self.accum

        self.step += 1
        self.beta_t *= self.beta
        return V


class AddSign(Optimizer):
    # paper: https://arxiv.org/abs/1709.07417

    def __init__(self, lr=0.01, mu=0.9, alpha=1):
        self.mu = _f(mu)
        self.alpha = _f(alpha)

        super().__init__(lr)

    def reset(self):
        self.accum = None

    def compute(self, dW, W):
        if self.accum is None:
            self.accum = np.zeros_like(dW)

        self.accum[:] = self.accum * self.mu + dW

        signed = np.sign(dW) * np.sign(self.accum)
        # signed *= decay

        return -self.lr * dW * (self.alpha + signed)


class PowerSign(Optimizer):
    # paper: https://arxiv.org/abs/1709.07417

    def __init__(self, lr=0.01, mu=0.9, alpha=np.e):
        self.mu = _f(mu)
        self.alpha = _f(alpha)
        self.use_exp = np.isclose(self.alpha, _f(np.e))

        super().__init__(lr)

    def reset(self):
        self.accum = None

    def compute(self, dW, W):
        if self.accum is None:
            self.accum = np.zeros_like(dW)

        self.accum[:] = self.accum * self.mu + dW

        signed = np.sign(dW) * np.sign(self.accum)
        # signed *= decay

        if self.use_exp:
            return -self.lr * dW * np.exp(signed)
        else:
            return -self.lr * dW * np.power(self.alpha, signed)


class Neumann(Optimizer):
    # paper: https://arxiv.org/abs/1712.03298
    # NOTE: this implementation is missing resetting as described in the paper.
    #       resetting is totally disabled for now.
    # NOTE: this implementation does not use vanilla SGD for its first epochs.
    #       you should do this yourself if you need it.
    #       it seems like using a Learner like SineCLR makes this unnecessary.

    def __init__(self, lr=0.01):
        self.alpha = _f(1e-7)  # cubic.
        self.beta = _f(1e-5)  # repulsive. NOTE: multiplied by len(dW) later.
        self.gamma = _f(0.99)  # EMA, or 1-pole low-pass parameter. same thing.
        # momentum is ∝ (in the shape of) 1 - 1/(1 + t)
        self.mu_min = _f(0.5)
        self.mu_max = _f(0.9)
        self.reset_period = 0  # TODO

        super().__init__(lr)

    def reset(self):
        # NOTE: mt and vt are different than the pair in Adam-like optimizers.
        self.mt = None  # momentum accumulator.
        self.vt = None  # weight accumulator.
        self.t = 0

    def compute(self, dW, W):
        raise Exception("compute() is not available for this Optimizer.")

    def update(self, dW, W):
        self.t += 1

        if self.mt is None:
            self.mt = np.zeros_like(dW)
        if self.vt is None:
            self.vt = np.zeros_like(dW)

        if self.reset_period > 0 and (self.t - 1) % self.reset_period == 0:
            self.mt = -self.lr * dW
            return

        # momentum quantity:
        mu = _1 - _1/_f(self.t)  # the + 1 is implicit.
        mu = (mu + self.mu_min) * (self.mu_max - self.mu_min)

        # smoothed change in weights:
        delta = W - self.vt
        delta_norm_squared = np.square(delta).sum()
        delta_norm = np.sqrt(delta_norm_squared)

        # regularization terms: (push and pull)
        cubic_reg = self.alpha * delta_norm_squared
        repulsive_reg = self.beta * dW.size / delta_norm_squared
        dt = dW + (cubic_reg - repulsive_reg) * (delta / delta_norm)

        # plain momentum:
        self.mt = mu * self.mt - self.lr * dt

        # weights and accumulator:
        W += mu * self.mt - self.lr * dt
        self.vt = W + self.gamma * (self.vt - W)


class AMSgrad(Optimizer):
    # paper: https://openreview.net/forum?id=ryQu7f-RZ
    # based on Adam. this simply adds a running element-wise maximum to vt.

    def __init__(self, lr=0.002, b1=0.9, b2=0.999, eps=1e-8, debias=True):
        self.b1 = _f(b1)  # decay term
        self.b2 = _f(b2)  # decay term
        self.b1_t_default = _f(b1)  # decay term power t
        self.b2_t_default = _f(b2)  # decay term power t
        self.eps = _f(eps)
        self.debias = bool(debias)

        super().__init__(lr)

    def reset(self):
        self.mt = None
        self.vt = None
        self.vtmax = None
        self.b1_t = self.b1_t_default
        self.b2_t = self.b2_t_default

    def compute(self, dW, W):
        if self.mt is None:
            self.mt = np.zeros_like(dW)
        if self.vt is None:
            self.vt = np.zeros_like(dW)
        if self.vtmax is None:
            self.vtmax = np.zeros_like(dW)

        # filter
        self.mt += (1 - self.b1) * (dW - self.mt)
        self.vt += (1 - self.b2) * (np.square(dW) - self.vt)

        self.vtmax = np.maximum(self.vtmax, self.vt)

        if self.debias:
            ret = -self.lr * (self.mt / (1 - self.b1_t)) \
                / (np.sqrt(self.vtmax / (1 - self.b2_t)) + self.eps)
        else:
            ret = -self.lr * self.mt / (np.sqrt(self.vtmax) + self.eps)

        # decay gain
        self.b1_t *= self.b1
        self.b2_t *= self.b2

        return ret