import numpy as np
from numpy import sin, cos, arcsin, arccos, arctan2
from ellipsoide import EllipsoidTriaxial
import matplotlib.pyplot as plt
from panou import louville_constant, func_sigma_ell, gha1_ana
import plotly.graph_objects as go
import winkelumrechnungen as wu

def gha1(ell: EllipsoidTriaxial, p0: np.ndarray, alpha0: float, s: float, ds: int):
    l0 = louville_constant(ell, p0, alpha0)
    points = [p0]
    alphas = [alpha0]
    s_curr = 0.0
    while s_curr < s:
        ds_step = min(ds, s - s_curr)
        if ds_step < 1e-8:
            break
        p1 = points[-1]
        alpha1 = alphas[-1]
        x1, y1, z1 = p1
        sigma = func_sigma_ell(ell, x1, y1, z1, alpha1)
        p2 = p1 + ds_step * sigma
        p2, _, _, _ = ell.cartonell(p2)
        ds_step = np.linalg.norm(p2 - p1)

        points.append(p2)
        dalpha = 1e-6
        l2 = louville_constant(ell, p2, alpha1)
        dl_dalpha = (louville_constant(ell, p2, alpha1+dalpha) - l2) / dalpha
        alpha2 = alpha1 + (l0 - l2) / dl_dalpha
        alphas.append(alpha2)
        s_curr += ds_step
    return points[-1], alphas[-1], np.array(points)

def show_points(points, p0, p1):
    fig = go.Figure()

    fig.add_scatter3d(x=points[:, 0], y=points[:, 1], z=points[:, 2],
                      mode='lines', line=dict(color="red", width=3), name="Approx")
    fig.add_scatter3d(x=[p0[0]], y=[p0[1]], z=[p0[2]],
                      mode='markers', marker=dict(color="green"), name="P0")
    fig.add_scatter3d(x=[p1[0]], y=[p1[1]], z=[p1[2]],
                      mode='markers', marker=dict(color="green"), name="P1")

    fig.update_layout(
        scene=dict(xaxis_title='X [km]',
                   yaxis_title='Y [km]',
                   zaxis_title='Z [km]',
                   aspectmode='data'),
        title="CHAMP")

    fig.show()


if __name__ == '__main__':
    ell = EllipsoidTriaxial.init_name("BursaSima1980round")
    P0 = ell.para2cart(0, 0)
    alpha0 = wu.deg2rad(90)
    s = 1000000
    P1_ana, alpha1_ana = gha1_ana(ell, P0, alpha0, s, 60, maxPartCircum=32)
    P1_app, alpha1_app, points = gha1(ell, P0, alpha0, s, 5000)
    show_points(points, P0, P1_ana)
    print(np.linalg.norm(P1_app - P1_ana))