import numpy as np
from numpy import arccos
from Hansen_ES_CMA import escma
from ellipsoide import EllipsoidTriaxial
from numpy.typing import NDArray
import plotly.graph_objects as go
from GHA_triaxial.panou_2013_2GHA_num import gha2_num
from utils import sigma2alpha

# Globals für Fitness
ell_ES: EllipsoidTriaxial = None
P_left: NDArray = None
P_right: NDArray = None


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 midpoint_fitness(x: tuple) -> float:
    """
    Fitness für einen Mittelpunkt P_middle zwischen P_left und P_right auf dem Ellipsoid:
    - Minimiert d(P_left, P_middle) + d(P_middle, P_right)
    - Erzwingt (weich) d(P_left,P_middle) ≈ d(P_middle,P_right)  (echter Mittelpunkt im Sinne der Polygonkette)
    """
    global ell_ES, P_left, P_right

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

    # midpoint penalty (dimensionslos)
    # relative Differenz, skaliert stabil ü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


def gha2_ES(ell: EllipsoidTriaxial, P0: NDArray, Pk: NDArray, maxSegLen: float = None, stopeval: int = 2000, maxIter: int = 10000, all_points: bool = False):
    """
    - Start mit [P0, Pk]
    - Für jedes Segment > maxSegLen: Mittelpunkt per CMA-ES optimieren und einfügen
    - Wiederholen bis alle Segmentlängen <= maxSegLen sind
    """
    global ell_ES
    ell_ES = ell
    R0 = (ell.ax + ell.ay + ell.b) / 3
    if maxSegLen is None:
        maxSegLen = R0 * 1 / (637.4) # 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_ES.cart2para(A)
                Bu, Bv = ell_ES.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, stopfitness=-np.inf,
                             stopeval=stopeval)

                P_next = ell.para2cart(u, v)
                new_points.append(P_next)
                startIter += 1
                if startIter > maxIter:
                    raise RuntimeError("Abbruch: 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_ES.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_ES, sigma0, P0)

        p1i = ell_ES.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_ES, 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.7, 0.3)
    P1 = ell.ell2cart(beta1, lamb1)

    alpha0, alpha1, s_num, betas, lambs = gha2_num(ell, beta0, lamb0, beta1, lamb1, n=10000, 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, all_points=True)
    print(s_num)
    print(s)
    print(s - s_num)
    show_points(points, points_num, P0, P1)