import numpy as np
from numpy import sin, cos, arcsin, arccos, arctan2
from ellipsoide import EllipsoidTriaxial
import matplotlib.pyplot as plt

dbeta_dc = lambda ell, beta, lamb, alpha: -2 * ell.Ey**2 * sin(alpha)**2 * cos(beta) * sin(beta)
dlamb_dc = lambda ell, beta, lamb, alpha: -2 * ell.Ee**2 * cos(alpha)**2 * sin(lamb) * cos(lamb)
dalpha_dc = lambda ell, beta, lamb, alpha:  (2 * sin(alpha) * cos(alpha) *
                                             (ell.Ey**2 * cos(beta)**2 + ell.Ee**2 * sin(lamb)**2))

lamb2_sphere = lambda r, phi1, lamb1, a12, s: lamb1 + arctan2(sin(s/r) * sin(a12),
                                                              cos(phi1) * cos(s/r) - sin(s/r) * cos(a12))
phi2_sphere = lambda r, phi1, lamb1, a12, s: arcsin(sin(phi1) * cos(s/r) + cos(phi1) * sin(s/r) * cos(a12))

a12_sphere = lambda phi1, lamb1, phi2, lamb2: arctan2(cos(phi2) * sin(lamb2 - lamb1),
                                                      cos(phi1) * cos(phi2) -
                                                      sin(phi1) * cos(phi2) * cos(lamb2 - lamb1))
s_sphere = lambda r, phi1, lamb1, phi2, lamb2: r * arccos(sin(phi1) * sin(phi2) +
                                                          cos(phi1) * cos(phi2) * cos(lamb2 - lamb1))

louville = lambda beta, lamb, alpha: (ell.Ey**2 * cos(beta)**2 * sin(alpha)**2 -
                                          ell.Ee**2 * sin(lamb)**2 * cos(alpha)**2)

def points_approx_gha2(r: float, phi1: np.ndarray, lamb1: np.ndarray, phi2: np.ndarray, lamb2: np.ndarray, num: int = None, step_size: float = 10000):
    s_approx = s_sphere(r, phi1, lamb1, phi2, lamb2)
    if num is not None:
        step_size = s_approx / (num+1)
    a_approx = a12_sphere(phi1, lamb1, phi2, lamb2)

    points = [np.array([phi1, lamb1, a_approx])]
    current_s = step_size
    while current_s < s_approx:
        phi_n = phi2_sphere(r, phi1, lamb1, a_approx, current_s)
        lamb_n = lamb2_sphere(r, phi1, lamb1, a_approx, current_s)
        points.append(np.array([phi_n, lamb_n, a_approx]))
        current_s += step_size
    points.append(np.array([phi2, lamb2, a_approx]))
    return points


def num_update(ell: EllipsoidTriaxial, points, diffs):
    for i, (beta, lamb, alpha) in enumerate(points):
        dalpha = dalpha_dc(ell, beta, lamb, alpha)
        if i == 0 or i == len(points) - 1:
            grad = np.array([0, 0, dalpha])
        else:
            dbeta = dbeta_dc(ell, beta, lamb, alpha)
            dlamb = dlamb_dc(ell, beta, lamb, alpha)
            grad = np.array([dbeta, dlamb, dalpha])

        delta = -diffs[i] * grad / np.dot(grad, grad)
        points[i] += delta
    return points


def gha2(ell: EllipsoidTriaxial, p1: np.ndarray, p2: np.ndarray, maxI: int):
    beta1, lamb1 = ell.cart2ell(p1)
    beta2, lamb2 = ell.cart2ell(p2)
    points = points_approx_gha2(ell.ax, beta1, lamb1, beta2, lamb2, 5)

    for j in range(maxI):
        constants = [louville(point[0], point[1], point[2]) for point in points]
        mean_constant = np.mean(constants)
        diffs = constants - mean_constant
        if np.mean(np.abs(diffs)) > 10:
            points = num_update(ell, points, diffs)
        else:
            break
        for k in range(maxI):


            last_diff_alpha = points[-2][-1] - points[-3][-1]
            alpha_extrap = points[-2][-1] + last_diff_alpha
            if abs(alpha_extrap - points[-1][-1]) > 0.0005:
                pass
            else:
                break
            pass
    pass
    return points

def show_points(ell: EllipsoidTriaxial, points):
    points_cart = []
    for point in points:
        points_cart.append(ell.ell2cart(point[0], point[1]))
    points_cart = np.array(points_cart)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    ax.plot(points_cart[:, 0], points_cart[:, 1], points_cart[:, 2])
    ax.scatter(points_cart[:, 0], points_cart[:, 1], points_cart[:, 2])

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.show()


if __name__ == "__main__":
    ell = EllipsoidTriaxial.init_name("Eitschberger1978")
    p1 = np.array([4189000, 812000, 4735000])
    p2 = np.array([4090000, 868000, 4808000])
    p1, phi1, lamb1, h1 = ell.cartonell(p1)
    p2, phi2, lamb2, h2 = ell.cartonell(p2)
    points = gha2(ell, p1, p2, 10)
    show_points(ell, points)