from typing import List, Tuple
import numpy as np
from autograd import grad


class descent:
    def __init__(
            self,
            step: float = 0.1,
            max_iter: int = 50,
            convergence: float = 1e-3,
            initial_points: Tuple[float, float] = (-1, -0.9),
    ):
        self.step = step
        self.max_iter = max_iter
        self.convergence = convergence
        self.initial_points = initial_points
        self.dx = 1e-6

    def gradient_descent(self, test_function) -> float:
        initial_points = self.initial_points
        x_n_minus_one = initial_points[0]
        x_n = initial_points[1]
        a = 0
        while a < self.max_iter and abs(test_function(x_n) - test_function(x_n_minus_one)) > self.convergence:
            x_new = x_n - self.step * grad(test_function)(x_n)
            x_n_minus_one = x_n
            x_n = x_new
            a += 1
        return x_n

    def newtons_method(self, test_function) -> float:
        initial_points = self.initial_points
        x_n_minus_one = initial_points[0]
        x_n = initial_points[1]
        a = 0
        while a < self.max_iter and abs(test_function(x_n) - test_function(x_n_minus_one)) > self.convergence:
            x_new = x_n - \
                test_function(
                    x_n)/((test_function(x_n + self.dx) - test_function(x_n))/self.dx)
            x_n_minus_one = x_n
            x_n = x_new
            a += 1
        return x_n

    def newtons_method_minimum(self, test_function) -> float:
        initial_points = self.initial_points
        x_n_minus_one = initial_points[0]
        x_n = initial_points[1]
        a = 0
        while a < self.max_iter and abs(test_function(x_n) - test_function(x_n_minus_one)) > self.convergence:
            x_new = x_n - grad(test_function)(x_n) / \
                grad(grad(test_function))(x_n)
            x_n_minus_one = x_n
            x_n = x_new
            a += 1
        return x_n

    def backtracking_line_search(
            self,
            test_function,
            current_point: float,
            search_direction: List[float],
            alpha: float = 0.2,
            beta: float = 0.9,
    ) -> float:
        full_step = 1
        p = search_direction
        x = current_point
        while test_function(x + full_step * p) > test_function(x) + alpha * full_step * np.dot(grad(test_function)(x), p):
            full_step *= beta
        return full_step

    def BFGS(self, test_function) -> float:
        initial_points = self.initial_points
        x_n_minus_one = initial_points[0]
        x_n = initial_points[1]
        a = 0
        Hessian_k = 1
        grad_k = grad(test_function)(x_n)
        while a < self.max_iter and abs(grad_k) > self.convergence:
            p_k = -np.dot(Hessian_k, grad(test_function)(x_n))
            alpha_k = self.backtracking_line_search(test_function, x_n, p_k)
            x_new = x_n + alpha_k * p_k
            grad_k = grad(test_function)(x_new)
            delta_x_k = x_new - x_n
            delta_g_k = grad_k - grad(test_function)(x_n)
            Hessian_k = Hessian_k + (1 + (np.dot(np.dot(Hessian_k, grad_k), grad_k)) / (np.dot(grad_k, p_k))) * np.dot(p_k, p_k.T) / np.dot(p_k, grad_k) \
                - (np.dot(np.dot(p_k, grad_k.T), Hessian_k) + np.dot(Hessian_k,
                   grad_k) * np.dot(p_k, grad_k.T)) / (np.dot(grad_k, p_k))
        return x_n

    def run_all(self, test_function) -> List[float]:
        return [self.gradient_descent(test_function), self.newtons_method(test_function), self.newtons_method_minimum(test_function), self.BFGS(test_function)]