from typing import Tuple

import numpy as np
import plotly.graph_objects as go
from numpy import cos, sin
from numpy.typing import NDArray

import winkelumrechnungen as wu
from GHA_triaxial.utils import louville_constant, pq_ell
from ellipsoid_triaxial import EllipsoidTriaxial
from utils_angle import wrap_0_2pi


def gha1_approx(ell: EllipsoidTriaxial, p0: np.ndarray, alpha0: float, s: float, ds: float, all_points: bool = False) \
        -> Tuple[NDArray, float] | Tuple[NDArray, float, NDArray, NDArray]:
    """
    Berechung einer Näherungslösung der ersten Hauptaufgabe
    :param ell: Ellipsoid
    :param p0: Anfangspunkt
    :param alpha0: Azimut im Anfangspunkt
    :param s: Strecke bis zum Endpunkt
    :param ds: Länge einzelner Streckenelemente
    :param all_points: Ausgabe aller Punkte als Array?
    :return: Endpunkt, Azimut im Endpunkt, optional alle Punkte
    """
    l0 = louville_constant(ell, p0, alpha0)
    points = [p0]
    alphas = [alpha0]
    s_curr = 0.0
    last_sigma = None
    last_p = None

    while s_curr < s:
        ds_step = min(ds, s - s_curr)
        if ds_step < 1e-8:
            break

        p1 = points[-1]
        alpha1 = alphas[-1]

        p, q = pq_ell(ell, p1)
        if last_p is not None and np.dot(p, last_p) < 0:
            p = -p
            q = -q
        last_p = p
        sigma = p * sin(alpha1) + q * cos(alpha1)
        if last_sigma is not None and np.dot(sigma, last_sigma) < 0:
            sigma = -sigma
            alpha1 += np.pi
            alpha1 = wrap_0_2pi(alpha1)
        p2 = p1 + ds_step * sigma
        p2 = ell.point_onto_ellipsoid(p2)

        dalpha = 1e-9
        l2 = louville_constant(ell, p2, alpha1)
        dl_dalpha = (louville_constant(ell, p2, alpha1+dalpha) - l2) / dalpha
        if abs(dl_dalpha) < 1e-20:
            alpha2 = alpha1 + 0
        else:
            alpha2 = alpha1 + (l0 - l2) / dl_dalpha
        points.append(p2)
        alphas.append(wrap_0_2pi(alpha2))

        ds_step = np.linalg.norm(p2 - p1)
        s_curr += ds_step
        last_sigma = sigma
        pass

    if all_points:
        return points[-1], alphas[-1], np.array(points), np.array(alphas)
    else:
        return points[-1], alphas[-1]

def show_points(points: NDArray, p0: NDArray, p1: NDArray):
    """
    Anzeigen der Punkte
    :param points: Array aller approximierten Punkte
    :param p0: Startpunkt
    :param p1: wahrer Endpunkt
    """
    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("KarneyTest2024")
    P0 = ell.ell2cart(wu.deg2rad(15), wu.deg2rad(15))
    alpha0 = wu.deg2rad(270)
    s = 1
    P1_app, alpha1_app, points, alphas = gha1_approx(ell, P0, alpha0, s, ds=0.1, all_points=True)
    # P1_ana, alpha1_ana = gha1_ana(ell, P0, alpha0, s, maxM=40, maxPartCircum=32)
    # print(np.linalg.norm(P1_app - P1_ana))
    # show_points(points, P0, P0)