import numpy as np
from Hansen_ES_CMA import escma
from ellipsoide import EllipsoidTriaxial
from numpy.typing import NDArray
from typing import Tuple
import plotly.graph_objects as go
from GHA_triaxial.gha2_num import gha2_num
from GHA_triaxial.utils import sigma2alpha, pq_ell





def Sehne(P1: NDArray, P2: NDArray) -> float:
    """
    Berechnung der 3D-Distanz zwischen zwei kartesischen Punkten
    :param P1: kartesische Koordinate Punkt 1
    :param P2: kartesische Koordinate Punkt 2
    :return: Bogenlänge s
    """
    R12 = P2-P1
    s = np.linalg.norm(R12)

    return s





def gha2_ES(ell: EllipsoidTriaxial, P0: NDArray, Pk: NDArray, maxSegLen: float = None, all_points: bool = False) -> Tuple[float, float, float, NDArray] | Tuple[float, float, float]:
    """
    Berechnen der 2. GHA mithilfe der CMA-ES.
    Die CMA-ES optimiert sukzessive den Mittelpunkt zwischen Start- und Zielpunkt. Der Abbruch der Berechnung erfolgt, wenn alle Segmentlängen <= maxSegLen sind.
    Die Distanzen zwischen den einzelnen Punkten werden als direkte 3D-Distanzen berechnet und aufaddiert.
    :param ell: Ellipsoid
    :param P0: Startpunkt
    :param Pk: Zielpunkt
    :param maxSegLen: maximale Segmentlänge
    :param all_points: Ergebnisliste mit allen Punkte, die wahlweise mit ausgegeben wird
    :return: Richtungswinkel in RAD des Start- und Zielpunktes und Gesamtlänge
    """
    P_left: NDArray = None
    P_right: NDArray = None

    def midpoint_fitness(x: tuple) -> float:
        """
        Fitness für einen Mittelpunkt P_middle zwischen P_left und P_right auf dem triaxialen Ellipsoid:
        - Minimiert d(P_left, P_middle) + d(P_middle, P_right)
        - Erzwingt d(P_left,P_middle) ≈ d(P_middle,P_right)  (echter Mittelpunkt im Sinne der Polygonkette)
        :param x: enthält die Startwerte von u und v
        :return: Fitnesswert (f)
        """
        nonlocal P_left, P_right, ell

        u, v = x
        P_middle = ell.para2cart(u, v)
        d1 = Sehne(P_left, P_middle)
        d2 = Sehne(P_middle, P_right)
        base = d1 + d2

        # midpoint penalty (dimensionslos)
        # relative Differenz, skaliert über verschiedene Segmentlängen
        denom = max(base, 1e-9)
        pen_equal = ((d1 - d2) / denom) ** 2
        w_equal = 10.0

        f = base + denom * w_equal * pen_equal

        return f

    R0 = (ell.ax + ell.ay + ell.b) / 3
    if maxSegLen is None:
        maxSegLen = R0 * 1 / (637.4*2) # 10km Segment bei mittleren Erdradius

    sigma_uv_nom = 1e-3 * (maxSegLen / R0)  # ~1e-5
    points: list[NDArray] = [P0, Pk]
    startIter = 0
    level = 0

    while True:
        seg_lens = [Sehne(points[i], points[i+1]) for i in range(len(points)-1)]
        max_len = max(seg_lens)
        if max_len <= maxSegLen:
            break

        level += 1
        new_points: list[NDArray] = [points[0]]
        for i in range(len(points) - 1):
            A = points[i]
            B = points[i+1]
            dAB = Sehne(A, B)
            # print(dAB)

            if dAB > maxSegLen:
                # global P_left, P_right
                P_left, P_right = A, B
                Au, Av = ell.cart2para(A)
                Bu, Bv = ell.cart2para(B)
                u0 = (Au + Bu) / 2
                v0 = Av + 0.5 * np.arctan2(np.sin(Bv - Av), np.cos(Bv - Av))
                xmean = [u0, v0]

                sigmaStep = sigma_uv_nom * (Sehne(A, B) / maxSegLen)

                u, v = escma(midpoint_fitness, N=2, xmean=xmean, sigma=sigmaStep) # Aufruf CMA-ES

                P_next = ell.para2cart(u, v)
                new_points.append(P_next)
                startIter += 1
                maxIter = 10000
                if startIter > maxIter:
                    raise RuntimeError("GHA2_ES: maximale Iterationen überschritten")
            new_points.append(B)

        points = new_points
        # print(f"[Level {level}] Punkte: {len(points)} | max Segment: {max_len:.3f} m")

    P_all = np.vstack(points)
    totalLen = float(np.sum(np.linalg.norm(P_all[1:] - P_all[:-1], axis=1)))

    if len(points) >= 3:
        p0i = ell.point_onto_ellipsoid(P0 + 10.0 * (points[1] - P0) / np.linalg.norm(points[1] - P0))
        sigma0 = (p0i - P0) / np.linalg.norm(p0i - P0)
        alpha0 = sigma2alpha(ell, sigma0, P0)

        p1i = ell.point_onto_ellipsoid(Pk - 10.0 * (Pk - points[-2]) / np.linalg.norm(Pk - points[-2]))
        sigma1 = (Pk - p1i) / np.linalg.norm(Pk - p1i)
        alpha1 = sigma2alpha(ell, sigma1, Pk)
    else:
        alpha0 = None
        alpha1 = None

    if all_points:
        return alpha0, alpha1, totalLen, P_all
    return alpha0, alpha1, totalLen


def show_points(points: NDArray, pointsES: NDArray, p0: NDArray, p1: NDArray):
    """
    Anzeigen der Punkte
    :param points: wahre Punkte der Linie
    :param pointsES: Punkte der Linie aus ES
    :param p0: wahrer Startpunkt
    :param p1: wahrer Endpunkt
    """
    fig = go.Figure()

    fig.add_scatter3d(x=pointsES[:, 0], y=pointsES[:, 1], z=pointsES[:, 2],
                      mode='lines', line=dict(color="green", width=3), name="Numerisch")
    fig.add_scatter3d(x=points[:, 0], y=points[:, 1], z=points[:, 2],
                      mode='lines', line=dict(color="red", width=3), name="ES")
    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("Bursa1970")

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

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

    alpha0, alpha1, s, points = gha2_ES(ell, P0, P1)
    print(s_num)
    print(s)
    print(alpha0)
    print(alpha1)
    print(s - s_num)
    show_points(points, points_num, P0, P1)
