import numpy as np
import time

class Optimizer:
    def __init__(self, l, alpha, h, seed=121212) -> None:
        self.l = l # number of directions
        self.alpha = alpha
        self.h = h
        self.rnd_state = np.random.RandomState(seed=seed)

    def _approximate_gradient(self, target, x, t):
        pass

    def optimize(self, f, x0, f_star, T = 100, reps = 5):
        fvals = np.empty((reps, T// (2 * self.l)))
        tms = np.empty((reps, T// (2 * self.l)))
        gnorm = np.empty((reps, T// (2 * self.l)))
        fx = f(x0) - f_star
        fvals[:, 0] = fx
        tms[:, 0] = 0.0
        for r in range(reps):
            x = x0.copy()
            for t in range(1, T // (2 * self.l)):
                it_time = time.time()
                g = self._approximate_gradient(f, x, t)
                x = x - self.alpha(t) * g
                it_time = time.time() - it_time
                fvals[r, t] = f(x) - f_star
                gnorm[r, t] = np.linalg.norm(g, ord=2)**2
                tms[r, t] = it_time
        gnorm[:, 0] = gnorm[:, 1]
        return gnorm, fvals, tms
    
    
    
class RandomCenterGaussian(Optimizer):

    def _approximate_gradient(self, f, x, t):
        P = self.rnd_state.randn(x.shape[0], self.l)
        h = self.h(t)
        g = np.zeros(x.shape[0])
        fx = f(x)
        for i in range(self.l):
            g = g + ((f(x + h * P[:, i]) - f(x - h * P[:, i]) )/(2*h)) * P[:, i]
        return g
        
    
    
class RandomCenteredSphere(Optimizer):

    def _approximate_gradient(self, f, x, t):
        P = self.rnd_state.randn(x.shape[0], self.l)
        P /= np.linalg.norm(P, axis=0, ord=2)
        h = self.h(t)
        g = np.zeros(x.shape[0])
        for i in range(self.l):
            g = g + (((f(x + h * P[:, i]) - f(x - h * P[:, i]))/(2*h)) * P[:, i])
        return g
    

class RandomCenteredStructured(Optimizer):

    def _approximate_gradient(self, f, x, t):
        A = self.rnd_state.randn(x.shape[0], self.l)
        P = np.linalg.qr(A, mode='reduced')[0]
        h = self.h(t)
        g = np.zeros(x.shape[0])
        for i in range(self.l):
            g = g + ((f(x + h * P[:, i]) - f(x - h * P[:, i]))/(2*h)) * P[:, i]
        return g
    
    
class RandomCenteredHouseholder(Optimizer):
    
    def __init__(self, d, l, alpha, h, seed=121212) -> None:
        super().__init__(l, alpha, h, seed)
        self.I = np.eye(d, l)

    def _approximate_gradient(self, f, x, t):
        v = self.rnd_state.randn(x.shape[0])
        v /= np.linalg.norm(v)
        P = self.I - 2 * np.outer(v, v[:self.l])# np.linalg.qr(A, mode='reduced')[0]
        h = self.h(t)
        g = np.zeros(x.shape[0])
        for i in range(self.l):
            g = g + ((f(x + h * P[:, i]) - f(x - h * P[:, i]))/(2*h)) * P[:, i]
        return g