import numpy as np
from ellipsoide import EllipsoidTriaxial
from GHA_triaxial.panou_2013_2GHA_num import gha2_num
import plotly.graph_objects as go
import winkelumrechnungen as wu
from numpy.typing import NDArray
from typing import Tuple

from utils import sigma2alpha


def gha2(ell: EllipsoidTriaxial, p0: NDArray, p1: NDArray, ds: float, all_points: bool = False) -> Tuple[float, float, float] | Tuple[float, float, float, NDArray]:
    """
    Numerische Approximation für die zweite Hauptaufgabe
    :param ell: Ellipsoid
    :param p0: Startpunkt
    :param p1: Endpunkt
    :param ds: maximales Streckenelement
    :param all_points: Alle Punkte ausgeben?
    :return:
    """
    points = np.array([p0, p1])
    while True:
        new_points = []

        for i in range(len(points)-1):
            new_points.append(points[i])

            pi = points[i] + 1/2 * (points[i+1] - points[i])
            pi = ell.cartonell(pi)

            new_points.append(pi)

        new_points.append(points[-1])
        points = np.array(new_points)

        elements = np.array([np.linalg.norm(points[i] - points[i+1]) for i in range(len(points)-1)])

        if np.average(elements) < ds:
            break

    p0i = ell.cartonell(p0 + ds/100 * (points[1] - p0) / np.linalg.norm(points[1] - p0))
    sigma0 = (p0i - p0) / np.linalg.norm(p0i - p0)
    alpha0 = sigma2alpha(ell, sigma0, p0)

    p1i = ell.cartonell(p1 - ds/100 * (p1 - points[-2]) / np.linalg.norm(p1 - points[-2]))
    sigma1 = (p1 - p1i) / np.linalg.norm(p1 - p1i)
    alpha1 = sigma2alpha(ell, sigma1, p1)

    s = np.sum(np.array([np.linalg.norm(points[i] - points[i+1]) for i in range(len(points)-1)]))

    if all_points:
        return alpha0, alpha1, s, np.array(points)
    else:
        return alpha0, alpha1, s

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

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

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

    fig.show()


if __name__ == '__main__':
    ell = EllipsoidTriaxial.init_name("KarneyTest2024")

    beta0, lamb0 = (0.2, 0.1)
    P0 = ell.ell2cart(beta0, lamb0)
    beta1, lamb1 = (0.7, 0.3)
    P1 = ell.ell2cart(beta1, lamb1)

    alpha0_app, alpha1_app, s_app, points = gha2(ell, P0, P1, ds=1e-4, all_points=True)

    alpha0, alpha1, s, betas, lambs = gha2_num(ell, beta0, lamb0, beta1, lamb1, n=5000, all_points=True)
    points_ana = []
    for beta, lamb in zip(betas, lambs):
        points_ana.append(ell.ell2cart(beta, lamb))
    points_ana = np.array(points_ana)

    show_points(points_ana, points, P0, P1)
    print(f"Differenz s: {s_app - s} m")
    print(f"Differenz alpha0: {wu.rad2deg(alpha0_app - alpha0)}°")
    print(f"Differenz alpha1: {wu.rad2deg(alpha1_app - alpha1)}°")