import numpy as np
import basin_of_attraction as basin
import trigonometric as trig
import scipy.optimize as opt
import aux


def minimize_gradient(A, z):
    """
    Returns the local minimizer found by the algorithm and the corresponding
    matrix X.

    Prints the value of local min and the corresponding local minimizer found
    by the gradient descent algorithm on the objective function of ||grad||_2.
    """

    dim, _ = A.shape
    init = np.random.uniform(-2, 2, size=dim)
    result = opt.minimize(
        fun=lambda theta: gradient_norm_func(A, theta), x0=init)
    theta = result.x
    opt_val = result.fun
    T = np.empty((dim, dim))
    for i in range(dim):
        T[i, :] = theta[i]
        T[i, :] = T[i, :] - theta
    X = np.cos(T)
    print("The local minimum is: {}".format(opt_val))
    print("The corresponding minimizer is: \n{}".format(X))

    return theta, X


def half_circle_init(dim):
    """
    Returns theta that generates positive X from within a half circle.
    """

    return np.random.uniform(-np.pi / 2, np.pi / 2, size=dim)


def full_circle_init(dim):
    """
    Returns theta that generates positive X from within a full circle.
    """

    return np.random.uniform(-np.pi, np.pi, size=dim)


def half_circle_mat(dim):
    """
    Returns positive X generated by function half_circle_init.
    """

    theta = half_circle_init(dim)
    T = np.empty((dim, dim))
    for i in range(dim):
        T[i, :] = theta[i]
        T[i, :] = T[i, :] - theta
    X = np.cos(T)
    return X


def full_circle_mat(dim):
    """
    Returns positive X generated by function full_circle_init.
    """

    theta = full_circle_init(dim)
    T = np.empty((dim, dim))
    for i in range(dim):
        T[i, :] = theta[i]
        T[i, :] = T[i, :] - theta
    X = np.cos(T)
    return X


def full_circle_rect(dim):
    """
    Returns the low rank rectangular matrix by function full_circle_init.
    """

    theta = full_circle_init(dim)
    return rect_from_theta(theta)


def mat_from_theta(theta):
    """
    Returns the matrix X from a certain theta vector.
    """

    dim = theta.shape[0]
    T = np.empty((dim, dim))
    for i in range(dim):
        T[i, :] = theta[i]
        T[i, :] = T[i, :] - theta
    X = np.cos(T)
    return X


def rect_from_theta(theta):
    """
    Returns the low rank rectangular matrix from theta.
    """

    theta = theta.reshape((-1, 1))
    return np.hstack((np.cos(theta), np.sin(theta)))


def gradient_norm_func(A, theta):
    """
    Returns the 2 norm of the gradient vector (more precisely,
    2 norm of 1/2 of the gradient vector).
    """

    return np.linalg.norm(trig.trig_grad(A, theta))


def evaluate_curvature(A, X):
    """
    Returns the Hessian given the observation A and a point X on the manifold.
    """

    dim, _ = A.shape
    one = np.ones(dim).reshape((-1, 1))
    hess = A * X - np.diag((A * X).dot(one))
    return hess


def largest_nonzero_curv(hess):
    """
    Returns the largest nonzero curvature for the given Hessian.
    """

    curv_arr = aux.sorted_eigenvalues(hess).ravel()
    curv = curv_arr[-1]
    if curv < 10 ** (-6):
        curv = curv_arr[-2]
    return curv


def smallest_nonzero_curv(hess):
    """
    Returns the smallest nonzero curvature for the given Hessian.
    """

    curv_arr = aux.sorted_eigenvalues(hess).ravel()
    curv = curv_arr[0]
    if curv > - 10 ** (-6):
        curv = curv_arr[1]
    return curv

if __name__ == '__main__':
    # # Find the global optimizer of function ||grad||_2 and check its curvature.
    # A, z = basin.get_observation(10, 1, 'positive-rows')
    # theta, X = minimize_gradient(A, z)
    # hessian = evaluate_curvature(A, X)
    # print("The smallest curvature is : {}".format(
    #     smallest_nonzero_curv(hessian)))

    # Compute the curvature of a randomly generated point with positive X (
    # generated by restraining the range of theta) ATTENTION: this gives the
    # curvature of the objective function, not the one for gradient function.
    while True:
        A, z = basin.get_observation(10, 1, 'positive-rows')
        dim, _ = A.shape
        for i in range(1000):
            theta = half_circle_init(dim)
            # theta = half_circle_init(dim)
            # print("The prediction X is: {}".format(mat_from_theta(theta)))

            hessian = - trig.trig_hess(A, theta)
            # hessian = - trig.trig_hess(A, np.zeros(dim))

            curv_small = smallest_nonzero_curv(hessian)
            curv_large = largest_nonzero_curv(hessian)
            print("The smallest nonzero curvature is: {}".format(
                np.real(curv_small)))
            print("The largest nonzero curvature is: {}".format(
                np.real(curv_large)))
            if curv_small > 10 ** (-6):
                exit(0)
            elif curv_large < 10 ** (-6):
                exit(0)
