import numpy as np
from ellipsoide import EllipsoidTriaxial
from runge_kutta import rk4, rk4_step, rk4_end, rk4_integral
import GHA_triaxial.numeric_examples_karney as ne_karney
import GHA_triaxial.numeric_examples_panou as ne_panou
import winkelumrechnungen as wu
from typing import Tuple
from numpy.typing import NDArray
import ausgaben as aus
from utils_angle import cot, arccot, wrap_to_pi


def norm_a(a):
    if a < 0.0:
        a += np.pi
    return a

def sph_azimuth(beta1, lam1, beta2, lam2):
    dlam = wrap_to_pi(lam2 - lam1)
    y = np.sin(dlam) * np.cos(beta2)
    x = np.cos(beta1) * np.sin(beta2) - np.sin(beta1) * np.cos(beta2) * np.cos(dlam)
    a = np.arctan2(y, x)
    if a < 0:
        a += 2 * np.pi
    return a

# Panou 2013
def gha2_num(
    ell: EllipsoidTriaxial,
    beta_0: float,
    lamb_0: float,
    beta_1: float,
    lamb_1: float,
    n: int = 16000,
    epsilon: float = 10**-12,
    iter_max: int = 30,
    all_points: bool = False,
) -> Tuple[float, float, float] | Tuple[float, float, float, NDArray, NDArray]:
    """
    :param ell: Ellipsoid
    :param beta_0: Beta Punkt 0
    :param lamb_0: Lambda Punkt 0
    :param beta_1: Beta Punkt 1
    :param lamb_1: Lambda Punkt 1
    :param n: Anzahl Schritte
    :param epsilon: Genauigkeit
    :param iter_max: Maximale Iterationen
    :param all_points: Ausgabe aller Punkte
    :return: Azimut Startpunkt, Azumit Zielpunkt, Strecke
    """

    # Berechnung Koeffizienten, Gaußschen Fundamentalgrößen 1. Ordnung sowie deren Ableitungen
    def BETA_LAMBDA(beta, lamb):
        BETA = (ell.ay ** 2 * np.sin(beta) ** 2 + ell.b ** 2 * np.cos(beta) ** 2) / (
            ell.Ex ** 2 - ell.Ey ** 2 * np.sin(beta) ** 2
        )
        LAMBDA = (ell.ax ** 2 * np.sin(lamb) ** 2 + ell.ay ** 2 * np.cos(lamb) ** 2) / (
            ell.Ex ** 2 - ell.Ee ** 2 * np.cos(lamb) ** 2
        )

        BETA_ = (ell.ax ** 2 * ell.Ey ** 2 * np.sin(2 * beta)) / (
            ell.Ex ** 2 - ell.Ey ** 2 * np.sin(beta) ** 2
        ) ** 2
        LAMBDA_ = -(ell.b ** 2 * ell.Ee ** 2 * np.sin(2 * lamb)) / (
            ell.Ex ** 2 - ell.Ee ** 2 * np.cos(lamb) ** 2
        ) ** 2

        BETA__ = (
            (2 * ell.ax ** 2 * ell.Ey ** 4 * np.sin(2 * beta) ** 2)
            / (ell.Ex ** 2 - ell.Ey ** 2 * np.sin(beta) ** 2) ** 3
            + (2 * ell.ax ** 2 * ell.Ey ** 2 * np.cos(2 * beta))
            / (ell.Ex ** 2 - ell.Ey ** 2 * np.sin(beta) ** 2) ** 2
        )
        LAMBDA__ = (
            (2 * ell.b ** 2 * ell.Ee ** 4 * np.sin(2 * lamb) ** 2)
            / (ell.Ex ** 2 - ell.Ee ** 2 * np.cos(lamb) ** 2) ** 3
            - (2 * ell.b ** 2 * ell.Ee ** 2 * np.sin(2 * lamb))
            / (ell.Ex ** 2 - ell.Ee ** 2 * np.cos(lamb) ** 2) ** 2
        )

        E = BETA * (ell.Ey ** 2 * np.cos(beta) ** 2 + ell.Ee ** 2 * np.sin(lamb) ** 2)
        G = LAMBDA * (ell.Ey ** 2 * np.cos(beta) ** 2 + ell.Ee ** 2 * np.sin(lamb) ** 2)

        E_beta = (
            BETA_ * (ell.Ey ** 2 * np.cos(beta) ** 2 + ell.Ee ** 2 * np.sin(lamb) ** 2)
            - BETA * ell.Ey ** 2 * np.sin(2 * beta)
        )
        E_lamb = BETA * ell.Ee ** 2 * np.sin(2 * lamb)

        G_beta = -LAMBDA * ell.Ey ** 2 * np.sin(2 * beta)
        G_lamb = (
            LAMBDA_ * (ell.Ey ** 2 * np.cos(beta) ** 2 + ell.Ee ** 2 * np.sin(lamb) ** 2)
            + LAMBDA * ell.Ee ** 2 * np.sin(2 * lamb)
        )

        E_beta_beta = (
            BETA__ * (ell.Ey ** 2 * np.cos(beta) ** 2 + ell.Ee ** 2 * np.sin(lamb) ** 2)
            - 2 * BETA_ * ell.Ey ** 2 * np.sin(2 * beta)
            - 2 * BETA * ell.Ey ** 2 * np.cos(2 * beta)
        )
        E_beta_lamb = BETA_ * ell.Ee ** 2 * np.sin(2 * lamb)
        E_lamb_lamb = 2 * BETA * ell.Ee ** 2 * np.cos(2 * lamb)

        G_beta_beta = -2 * LAMBDA * ell.Ey ** 2 * np.cos(2 * beta)
        G_beta_lamb = -LAMBDA_ * ell.Ey ** 2 * np.sin(2 * beta)
        G_lamb_lamb = (
            LAMBDA__ * (ell.Ey ** 2 * np.cos(beta) ** 2 + ell.Ee ** 2 * np.sin(lamb) ** 2)
            + 2 * LAMBDA_ * ell.Ee ** 2 * np.sin(2 * lamb)
            + 2 * LAMBDA * ell.Ee ** 2 * np.cos(2 * lamb)
        )

        return (
            BETA,
            LAMBDA,
            E,
            G,
            BETA_,
            LAMBDA_,
            BETA__,
            LAMBDA__,
            E_beta,
            E_lamb,
            G_beta,
            G_lamb,
            E_beta_beta,
            E_beta_lamb,
            E_lamb_lamb,
            G_beta_beta,
            G_beta_lamb,
            G_lamb_lamb,
        )

    # Berechnung der ODE Koeffizienten für Fall 1 (lambda_0 != lambda_1)
    def p_coef(beta, lamb):
        (
            BETA,
            LAMBDA,
            E,
            G,
            BETA_,
            LAMBDA_,
            BETA__,
            LAMBDA__,
            E_beta,
            E_lamb,
            G_beta,
            G_lamb,
            E_beta_beta,
            E_beta_lamb,
            E_lamb_lamb,
            G_beta_beta,
            G_beta_lamb,
            G_lamb_lamb,
        ) = BETA_LAMBDA(beta, lamb)

        p_3 = -0.5 * (E_lamb / G)
        p_2 = (G_beta / G) - 0.5 * (E_beta / E)
        p_1 = 0.5 * (G_lamb / G) - (E_lamb / E)
        p_0 = 0.5 * (G_beta / E)

        p_33 = -0.5 * ((E_beta_lamb * G - E_lamb * G_beta) / (G**2))
        p_22 = ((G * G_beta_beta - G_beta * G_beta) / (G**2)) - 0.5 * (
            (E * E_beta_beta - E_beta * E_beta) / (E**2)
        )
        p_11 = 0.5 * ((G * G_beta_lamb - G_beta * G_lamb) / (G**2)) - (
            (E * E_beta_lamb - E_beta * E_lamb) / (E**2)
        )
        p_00 = 0.5 * ((E * G_beta_beta - E_beta * G_beta) / (E**2))

        return (BETA, LAMBDA, E, G, p_3, p_2, p_1, p_0, p_33, p_22, p_11, p_00)

    # Berechnung der ODE Koeffizienten für Fall 2 (lambda_0 == lambda_1)
    def q_coef(beta, lamb):
        (
            BETA,
            LAMBDA,
            E,
            G,
            BETA_,
            LAMBDA_,
            BETA__,
            LAMBDA__,
            E_beta,
            E_lamb,
            G_beta,
            G_lamb,
            E_beta_beta,
            E_beta_lamb,
            E_lamb_lamb,
            G_beta_beta,
            G_beta_lamb,
            G_lamb_lamb,
        ) = BETA_LAMBDA(beta, lamb)

        q_3 = -0.5 * (G_beta / E)
        q_2 = (E_lamb / E) - 0.5 * (G_lamb / G)
        q_1 = 0.5 * (E_beta / E) - (G_beta / G)
        q_0 = 0.5 * (E_lamb / G)

        q_33 = -0.5 * ((E * G_beta_lamb - E_lamb * G_lamb) / (E**2))
        q_22 = ((E * E_lamb_lamb - E_lamb * E_lamb) / (E**2)) - 0.5 * (
            (G * G_lamb_lamb - G_lamb * G_lamb) / (G**2)
        )
        q_11 = 0.5 * ((E * E_beta_lamb - E_beta * E_lamb) / (E**2)) - (
            (G * G_beta_lamb - G_beta * G_lamb) / (G**2)
        )
        q_00 = 0.5 * ((E_lamb_lamb * G - E_lamb * G_lamb) / (G**2))

        return BETA, LAMBDA, E, G, q_3, q_2, q_1, q_0, q_33, q_22, q_11, q_00

    def integrand_lambda(lamb, y):
        beta = y[0]
        beta_p = y[1]
        (_, _, E, G, *_) = BETA_LAMBDA(beta, lamb)
        return np.sqrt(E * beta_p**2 + G)

    def integrand_beta(beta, y):
        lamb = y[0]
        lamb_p = y[1]
        (_, _, E, G, *_) = BETA_LAMBDA(beta, lamb)
        return np.sqrt(E + G * lamb_p**2)

    # Fall 1 (lambda_0 != lambda_1)
    if abs(lamb_1 - lamb_0) >= 1e-15:
        N = int(n)
        dlamb = float(lamb_1 - lamb_0)

        beta0 = float(beta_0)
        lamb0 = float(lamb_0)
        beta1 = float(beta_1)
        lamb1 = float(lamb_1)

        def ode_lamb(lamb, v):
            beta, beta_p, X3, X4 = v
            (_, _, _, _, p_3, p_2, p_1, p_0, p_33, p_22, p_11, p_00) = p_coef(beta, lamb)

            dbeta = beta_p
            dbeta_p = p_3 * beta_p**3 + p_2 * beta_p**2 + p_1 * beta_p + p_0
            dX3 = X4
            dX4 = (p_33 * beta_p**3 + p_22 * beta_p**2 + p_11 * beta_p + p_00) * X3 + (3*p_3*beta_p**2 + 2*p_2*beta_p + p_1) * X4
            return np.array([dbeta, dbeta_p, dX3, dX4], dtype=float)

        alpha0_sph = sph_azimuth(beta0, lamb0, beta1, lamb1)
        (_, _, E0, G0, *_) = BETA_LAMBDA(beta0, lamb0)
        beta_p0_sph = np.sqrt(G0 / E0) * cot(alpha0_sph)

        N_newton = min(N, 4000)

        def solve_newton(beta_p0_init: float):
            beta_p0 = float(beta_p0_init)
            for _ in range(iter_max):
                v0 = np.array([beta0, beta_p0, 0.0, 1.0], dtype=float)
                _, y_end = rk4_end(ode_lamb, lamb0, v0, dlamb, N_newton)

                beta_end, _, X3_end, _ = y_end
                delta = beta_end - beta1

                if abs(delta) < epsilon:
                    return True, beta_p0

                if abs(X3_end) < 1e-20:
                    return False, None

                step = delta / X3_end
                step = np.clip(step, -0.5, 0.5)
                beta_p0 -= step

            return False, None

        ok, beta_p0_sol = solve_newton(beta_p0_sph)

        if not ok:
            candidates = [-beta_p0_sph, 0.5 * beta_p0_sph, 2.0 * beta_p0_sph]
            N_quick = min(N, 2000)
            best = None
            for g in candidates:
                ok_g, sol = solve_newton(g)
                if not ok_g:
                    continue
                v0_g = np.array([beta0, sol, 0.0, 1.0], dtype=float)
                _, _, s_quick = rk4_integral(ode_lamb, lamb0, v0_g, dlamb, N_quick, integrand_lambda)
                if (best is None) or (s_quick < best[0]):
                    best = (s_quick, sol)
            if best is None:
                raise RuntimeError("GHA2_num: Keine Startwert-Variante konvergiert (lambda-Fall)")
            beta_p0_sol = best[1]

        beta_p0 = float(beta_p0_sol)
        v0_final = np.array([beta0, beta_p0, 0.0, 1.0], dtype=float)

        if all_points:
            lamb_list, states = rk4(ode_lamb, lamb0, v0_final, dlamb, N, False)
            lamb_arr = np.array(lamb_list, dtype=float)
            beta_arr = np.array([st[0] for st in states], dtype=float)
            beta_p_arr = np.array([st[1] for st in states], dtype=float)

            (_, _, E_start, G_start, *_) = BETA_LAMBDA(beta_arr[0], lamb_arr[0])
            (_, _, E_end,   G_end,   *_) = BETA_LAMBDA(beta_arr[-1], lamb_arr[-1])

            alpha_1 = norm_a(arccot(np.sqrt(E_start / G_start) * beta_p_arr[0]))
            alpha_2 = norm_a(arccot(np.sqrt(E_end   / G_end)   * beta_p_arr[-1]))

            # Distanz aus Arrays
            integrand = np.zeros(N + 1, dtype=float)
            for i in range(N + 1):
                (_, _, Ei, Gi, *_) = BETA_LAMBDA(beta_arr[i], lamb_arr[i])
                integrand[i] = np.sqrt(Ei * beta_p_arr[i]**2 + Gi)

            h = abs(dlamb) / N
            if N % 2 == 0:
                S = integrand[0] + integrand[-1] + 4.0*np.sum(integrand[1:-1:2]) + 2.0*np.sum(integrand[2:-1:2])
                s = h/3.0 * S
            else:
                s = np.trapz(integrand, dx=h)

            return float(alpha_1), float(alpha_2), float(s), beta_arr, lamb_arr

        _, y_end, s = rk4_integral(ode_lamb, lamb0, v0_final, dlamb, N, integrand_lambda)
        beta_end, beta_p_end, _, _ = y_end

        (_, _, E_start, G_start, *_) = BETA_LAMBDA(beta0, lamb0)
        (_, _, E_end,   G_end,   *_) = BETA_LAMBDA(beta1, lamb1)

        alpha_1 = norm_a(arccot(np.sqrt(E_start / G_start) * beta_p0))
        alpha_2 = norm_a(arccot(np.sqrt(E_end   / G_end)   * beta_p_end))

        return float(alpha_1), float(alpha_2), float(s)

    # Fall 2 (lambda_0 == lambda_1)
    N = int(n)
    dbeta = float(beta_1 - beta_0)

    if abs(dbeta) < 1e-15:
        if all_points:
            return 0.0, 0.0, 0.0, np.array([]), np.array([])
        return 0.0, 0.0, 0.0

    beta0 = float(beta_0)
    lamb0 = float(lamb_0)
    beta1 = float(beta_1)
    lamb1 = float(lamb_1)

    def ode_beta(beta, v):
        lamb, lamb_p, Y3, Y4 = v
        (_, _, _, _, q_3, q_2, q_1, q_0, q_33, q_22, q_11, q_00) = q_coef(beta, lamb)

        dlamb = lamb_p
        dlamb_p = q_3*lamb_p**3 + q_2*lamb_p**2 + q_1*lamb_p + q_0
        dY3 = Y4
        dY4 = (q_33*lamb_p**3 + q_22*lamb_p**2 + q_11*lamb_p + q_00)*Y3 + (3*q_3*lamb_p**2 + 2*q_2*lamb_p + q_1)*Y4
        return np.array([dlamb, dlamb_p, dY3, dY4], dtype=float)

    lamb_p0 = 0.0
    for _ in range(iter_max):
        v0 = np.array([lamb0, lamb_p0, 0.0, 1.0], dtype=float)
        _, y_end = rk4_end(ode_beta, beta0, v0, dbeta, N)

        lamb_end, _, Y3_end, _ = y_end
        delta = lamb_end - lamb1

        if abs(delta) < epsilon:
            break

        if abs(Y3_end) < 1e-20:
            raise RuntimeError("GHA2_num: Ableitung ~ 0 im beta-Fall")

        step = delta / Y3_end
        step = np.clip(step, -1.0, 1.0)
        lamb_p0 -= step

    v0_final = np.array([lamb0, lamb_p0, 0.0, 1.0], dtype=float)

    if all_points:
        beta_list, states = rk4(ode_beta, beta0, v0_final, dbeta, N, False)
        beta_arr = np.array(beta_list, dtype=float)
        lamb_arr = np.array([st[0] for st in states], dtype=float)
        lamb_p_arr = np.array([st[1] for st in states], dtype=float)

        (BETA_s, LAMBDA_s, _, _, *_) = BETA_LAMBDA(beta_arr[0],  lamb_arr[0])
        (BETA_e, LAMBDA_e, _, _, *_) = BETA_LAMBDA(beta_arr[-1], lamb_arr[-1])

        alpha_1 = norm_a((np.pi/2.0) - arccot(np.sqrt(LAMBDA_s / BETA_s) * lamb_p_arr[0]))
        alpha_2 = norm_a((np.pi/2.0) - arccot(np.sqrt(LAMBDA_e / BETA_e) * lamb_p_arr[-1]))

        integrand = np.zeros(N + 1, dtype=float)
        for i in range(N + 1):
            (_, _, Ei, Gi, *_) = BETA_LAMBDA(beta_arr[i], lamb_arr[i])
            integrand[i] = np.sqrt(Ei + Gi * lamb_p_arr[i]**2)

        h = abs(dbeta) / N
        if N % 2 == 0:
            S = integrand[0] + integrand[-1] + 4.0*np.sum(integrand[1:-1:2]) + 2.0*np.sum(integrand[2:-1:2])
            s = h/3.0 * S
        else:
            s = np.trapz(integrand, dx=h)

        return float(alpha_1), float(alpha_2), float(s), beta_arr, lamb_arr

    _, y_end, s = rk4_integral(ode_beta, beta0, v0_final, dbeta, N, integrand_beta)
    lamb_end, lamb_p_end, _, _ = y_end

    (BETA_s, LAMBDA_s, _, _, *_) = BETA_LAMBDA(beta0, lamb0)
    (BETA_e, LAMBDA_e, _, _, *_) = BETA_LAMBDA(beta1, lamb1)

    alpha_1 = norm_a((np.pi/2.0) - arccot(np.sqrt(LAMBDA_s / BETA_s) * lamb_p0))
    alpha_2 = norm_a((np.pi/2.0) - arccot(np.sqrt(LAMBDA_e / BETA_e) * lamb_p_end))

    return float(alpha_1), float(alpha_2), float(s)

if __name__ == "__main__":
    # ell = EllipsoidTriaxial.init_name("BursaSima1980round")
    # beta1 = np.deg2rad(75)
    # lamb1 = np.deg2rad(-90)
    # beta2 = np.deg2rad(75)
    # lamb2 = np.deg2rad(66)
    # a0, a1, s = gha2_num(ell, beta1, lamb1, beta2, lamb2, n=5000)
    # print(aus.gms("a0", a0, 4))
    # print(aus.gms("a1", a1, 4))
    # print("s: ", s)
    # # print(aus.gms("a2", a2, 4))
    # # print(s)
    # cart1 = ell.para2cart(0, 0)
    # cart2 = ell.para2cart(0.4, 1.4)
    # beta1, lamb1 = ell.cart2ell(cart1)
    # beta2, lamb2 = ell.cart2ell(cart2)
    #
    # a1, a2, s = gha2_num(ell, beta1, lamb1, beta2, lamb2, n=5000)
    # print(s)

    # ell = EllipsoidTriaxial.init_name("BursaSima1980round")
    # diffs_panou = []
    # examples_panou = ne_panou.get_random_examples(4)
    # for example in examples_panou:
    #     beta0, lamb0, beta1, lamb1, _, alpha0, alpha1, s = example
    #     P0 = ell.ell2cart(beta0, lamb0)
    #     try:
    #         alpha0_num, alpha1_num, s_num = gha2_num(ell, beta0, lamb0, beta1, lamb1, n=4000, iter_max=10)
    #         diffs_panou.append(
    #             (wu.rad2deg(abs(alpha0 - alpha0_num)), wu.rad2deg(abs(alpha1 - alpha1_num)), abs(s - s_num)))
    #     except:
    #         print(f"Fehler für {beta0}, {lamb0}, {beta1}, {lamb1}")
    # diffs_panou = np.array(diffs_panou)
    # print(diffs_panou)
    #
    # ell = EllipsoidTriaxial.init_name("KarneyTest2024")
    # diffs_karney = []
    # # examples_karney = ne_karney.get_examples((30500, 40500))
    # examples_karney = ne_karney.get_random_examples(2)
    # for example in examples_karney:
    #     beta0, lamb0, alpha0, beta1, lamb1, alpha1, s = example
    #
    #     try:
    #         alpha0_num, alpha1_num, s_num = gha2_num(ell, beta0, lamb0, beta1, lamb1, n=4000, iter_max=10)
    #         diffs_karney.append((wu.rad2deg(abs(alpha0-alpha0_num)), wu.rad2deg(abs(alpha1-alpha1_num)), abs(s-s_num)))
    #     except:
    #         print(f"Fehler für {beta0}, {lamb0}, {beta1}, {lamb1}")
    # diffs_karney = np.array(diffs_karney)
    # print(diffs_karney)

    pass