import time
import pickle
import numpy as np
from numpy import nan
import winkelumrechnungen as wu
import os
from contextlib import contextmanager, redirect_stdout, redirect_stderr

from ellipsoide import EllipsoidTriaxial
from GHA_triaxial.panou import alpha_ell2para, alpha_para2ell

from GHA_triaxial.panou import gha1_num, gha1_ana
from GHA_triaxial.approx_gha1 import gha1_approx

from GHA_triaxial.panou_2013_2GHA_num import gha2_num
from GHA_triaxial.ES_gha2 import gha2_ES
from GHA_triaxial.approx_gha2 import gha2_approx

from GHA_triaxial.numeric_examples_panou import get_random_examples as get_examples_panou
from GHA_triaxial.numeric_examples_karney import get_random_examples as get_examples_karney

@contextmanager
def suppress_print():
    with open(os.devnull, 'w') as fnull:
        with redirect_stdout(fnull), redirect_stderr(fnull):
            yield

# steps_gha1_num = [2000, 5000, 10000, 20000]
# maxM_gha1_ana = [20, 40, 60]
# parts_gha1_ana = [4, 8, 16]
# dsPart_gha1_approx = [600, 1250, 6000, 60000]   # entspricht bei der Erde ca. 10000, 5000, 1000, 100
#
# steps_gha2_num = [2000, 5000, 10000, 20000]
# dsPart_gha2_ES = [600, 1250, 6000]   # entspricht bei der Erde ca. 10000, 5000, 1000
# dsPart_gha2_approx = [600, 1250, 6000, 60000]   # entspricht bei der Erde ca. 10000, 5000, 1000, 100

steps_gha1_num = [2000, 5000]
maxM_gha1_ana = [20, 40]
parts_gha1_ana = [4, 8]
dsPart_gha1_approx = [600, 1250]

steps_gha2_num = [2000, 5000]
dsPart_gha2_ES = [20]
dsPart_gha2_approx = [600, 1250]

ell_karney: EllipsoidTriaxial = EllipsoidTriaxial.init_name("KarneyTest2024")
ell_panou: EllipsoidTriaxial = EllipsoidTriaxial.init_name("BursaSima1980round")

results_karney = {}
results_panou = {}

examples_karney = get_examples_karney(2, 42)
examples_panou = get_examples_panou(2, 42)

for example in examples_karney:
    example_results = {}

    beta0, lamb0, alpha0_ell, beta1, lamb1, alpha1_ell, s = example
    P0 = ell_karney.ell2cart(beta0, lamb0)
    P1 = ell_karney.ell2cart(beta1, lamb1)
    _, _, alpha0_para = alpha_ell2para(ell_karney, beta0, lamb0, alpha0_ell)

    for steps in steps_gha1_num:
        start = time.perf_counter()
        try:
            P1_num, alpha1_num_1 = gha1_num(ell_karney, P0, alpha0_ell, s, num=steps)
            end = time.perf_counter()
            beta1_num, lamb1_num = ell_karney.cart2ell(P1_num)
            d_beta1 = abs(wu.rad2deg(beta1_num - beta1)) / 3600
            d_lamb1 = abs(wu.rad2deg(lamb1_num - lamb1)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_num_1 - alpha1_ell)) / 3600
            d_time = end - start
            example_results[f"GHA1_num_{steps}"] = (d_beta1, d_lamb1, d_alpha1, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA1_num_{steps}"] = (nan, nan, nan, nan)

    for maxM in maxM_gha1_ana:
        for parts in parts_gha1_ana:
            start = time.perf_counter()
            try:
                P1_ana, alpha1_ana_para = gha1_ana(ell_karney, P0, alpha0_para, s, maxM=maxM, maxPartCircum=parts)
                end = time.perf_counter()
                beta1_ana, lamb1_ana = ell_karney.cart2ell(P1_ana)
                _, _, alpha1_ana_ell = alpha_para2ell(ell_karney, beta1_ana, lamb1_ana, alpha1_ana_para)
                d_beta1 = abs(wu.rad2deg(beta1_ana - beta1)) / 3600
                d_lamb1 = abs(wu.rad2deg(lamb1_ana - lamb1)) / 3600
                d_alpha1 = abs(wu.rad2deg(alpha1_ana_ell - alpha1_ell)) / 3600
                d_time = end - start
                example_results[f"GHA1_ana_{maxM}_{parts}"] = (d_beta1, d_lamb1, d_alpha1, d_time)
            except Exception as e:
                print(e)
                example_results[f"GHA1_ana_{maxM}_{parts}"] = (nan, nan, nan, nan)

    for dsPart in dsPart_gha1_approx:
        ds = ell_karney.ax/dsPart
        start = time.perf_counter()
        try:
            P1_approx, alpha1_approx = gha1_approx(ell_karney, P0, alpha0_ell, s, ds=ds)
            end = time.perf_counter()
            beta1_approx, lamb1_approx = ell_karney.cart2ell(P1_approx)
            d_beta1 = abs(wu.rad2deg(beta1_approx - beta1)) / 3600
            d_lamb1 = abs(wu.rad2deg(lamb1_approx - lamb1)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_approx - alpha1_ell)) / 3600
            d_time = end - start
            example_results[f"GHA1_approx_{ds:.3f}"] = (d_beta1, d_lamb1, d_alpha1, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA1_approx_{ds:.3f}"] = (nan, nan, nan, nan)

    for steps in steps_gha2_num:
        start = time.perf_counter()
        try:
            alpha0_num, alpha1_num_2, s_num = gha2_num(ell_karney, beta0, lamb0, beta1, lamb1, n=steps)
            end = time.perf_counter()
            d_alpha0 = abs(wu.rad2deg(alpha0_num - alpha0_ell)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_num_2 - alpha1_ell)) / 3600
            d_s = abs(s_num - s) / 1000
            d_time = end - start
            example_results[f"GHA2_num_{steps}"] = (d_alpha0, d_alpha1, d_s, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA2_num_{steps}"] = (nan, nan, nan, nan)

    for dsPart in dsPart_gha2_ES:
        ds = ell_karney.ax/dsPart
        start = time.perf_counter()
        try:
            with suppress_print():
                alpha0_ES, alpha1_ES, s_ES = gha2_ES(ell_karney, P0, P1, stepLenTarget=ds)
            end = time.perf_counter()
            d_alpha0 = abs(wu.rad2deg(alpha0_ES - alpha0_ell)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_ES - alpha1_ell)) / 3600
            d_s = abs(s_ES - s) / 1000
            d_time = end - start
            example_results[f"GHA2_ES_{ds:.3f}"] = (d_alpha0, d_alpha1, d_s, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA2_ES_{ds:.3f}"] = (nan, nan, nan, nan)

    for dsPart in dsPart_gha2_approx:
        ds = ell_karney.ax/dsPart
        start = time.perf_counter()
        try:
            alpha0_approx, alpha1_approx, s_approx = gha2_approx(ell_karney, P0, P1, ds=ds)
            end = time.perf_counter()
            d_alpha0 = abs(wu.rad2deg(alpha0_approx - alpha0_ell)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_approx - alpha1_ell)) / 3600
            d_s = abs(s_approx - s) / 1000
            d_time = end - start
            example_results[f"GHA2_approx_{ds:.3f}"] = (d_alpha0, d_alpha1, d_s, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA2_approx_{ds:.3f}"] = (nan, nan, nan, nan)

    results_karney[f"beta0: {wu.rad2deg(beta0):.3f}, lamb0: {wu.rad2deg(lamb0):.3f}, alpha0: {wu.rad2deg(alpha0_ell):.3f}, s: {s}"] = example_results

for example in examples_panou:
    example_results = {}

    beta0, lamb0, beta1, lamb1, _, alpha0_ell, alpha1_ell, s = example
    P0 = ell_panou.ell2cart(beta0, lamb0)
    P1 = ell_panou.ell2cart(beta1, lamb1)
    _, _, alpha0_para = alpha_ell2para(ell_panou, beta0, lamb0, alpha0_ell)

    for steps in steps_gha1_num:
        start = time.perf_counter()
        try:
            P1_num, alpha1_num_1 = gha1_num(ell_panou, P0, alpha0_ell, s, num=steps)
            end = time.perf_counter()
            beta1_num, lamb1_num = ell_panou.cart2ell(P1_num)
            d_beta1 = abs(wu.rad2deg(beta1_num - beta1)) / 3600
            d_lamb1 = abs(wu.rad2deg(lamb1_num - lamb1)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_num_1 - alpha1_ell)) / 3600
            d_time = end - start
            example_results[f"GHA1_num_{steps}"] = (d_beta1, d_lamb1, d_alpha1, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA1_num_{steps}"] = (nan, nan, nan, nan)

    for maxM in maxM_gha1_ana:
        for parts in parts_gha1_ana:
            start = time.perf_counter()
            try:
                P1_ana, alpha1_ana_para = gha1_ana(ell_panou, P0, alpha0_para, s, maxM=maxM, maxPartCircum=parts)
                end = time.perf_counter()
                beta1_ana, lamb1_ana = ell_panou.cart2ell(P1_ana)
                _, _, alpha1_ana = alpha_para2ell(ell_panou, beta1_ana, lamb1_ana, alpha1_ana_para)
                d_beta1 = abs(wu.rad2deg(beta1_ana - beta1)) / 3600
                d_lamb1 = abs(wu.rad2deg(lamb1_ana - lamb1)) / 3600
                d_alpha1 = abs(wu.rad2deg(alpha1_ana - alpha1_ell)) / 3600
                d_time = end - start
                example_results[f"GHA1_ana_{maxM}_{parts}"] = (d_beta1, d_lamb1, d_alpha1, d_time)
            except Exception as e:
                print(e)
                example_results[f"GHA1_ana_{maxM}_{parts}"] = (nan, nan, nan, nan)

    for dsPart in dsPart_gha1_approx:
        ds = ell_panou.ax/dsPart
        start = time.perf_counter()
        try:
            P1_approx, alpha1_approx = gha1_approx(ell_panou, P0, alpha0_ell, s, ds=ds)
            end = time.perf_counter()
            beta1_approx, lamb1_approx = ell_panou.cart2ell(P1_approx)
            d_beta1 = abs(wu.rad2deg(beta1_approx - beta1)) / 3600
            d_lamb1 = abs(wu.rad2deg(lamb1_approx - lamb1)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_approx - alpha1_ell)) / 3600
            d_time = end - start
            example_results[f"GHA1_approx_{ds:.3f}"] = (d_beta1, d_lamb1, d_alpha1, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA1_approx_{ds:.3f}"] = (nan, nan, nan, nan)

    for steps in steps_gha2_num:
        start = time.perf_counter()
        try:
            alpha0_num, alpha1_num_2, s_num = gha2_num(ell_panou, beta0, lamb0, beta1, lamb1, n=steps)
            end = time.perf_counter()
            d_alpha0 = abs(wu.rad2deg(alpha0_num - alpha0_ell)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_num_2 - alpha1_ell)) / 3600
            d_s = abs(s_num - s) / 1000
            d_time = end - start
            example_results[f"GHA2_num_{steps}"] = (d_alpha0, d_alpha1, d_s, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA2_num_{steps}"] = (nan, nan, nan, nan)

    for dsPart in dsPart_gha2_ES:
        ds = ell_panou.ax/dsPart
        start = time.perf_counter()
        try:
            with suppress_print():
                alpha0_ES, alpha1_ES, s_ES = gha2_ES(ell_panou, P0, P1, stepLenTarget=ds)
            end = time.perf_counter()
            d_alpha0 = abs(wu.rad2deg(alpha0_ES - alpha0_ell)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_ES - alpha1_ell)) / 3600
            d_s = abs(s_ES - s) / 1000
            d_time = end - start
            example_results[f"GHA2_ES_{ds:.3f}"] = (d_alpha0, d_alpha1, d_s, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA2_ES_{ds:.3f}"] = (nan, nan, nan, nan)

    for dsPart in dsPart_gha2_approx:
        ds = ell_panou.ax/dsPart
        start = time.perf_counter()
        try:
            alpha0_approx, alpha1_approx, s_approx = gha2_approx(ell_panou, P0, P1, ds=ds)
            end = time.perf_counter()
            d_alpha0 = abs(wu.rad2deg(alpha0_approx - alpha0_ell)) / 3600
            d_alpha1 = abs(wu.rad2deg(alpha1_approx - alpha1_ell)) / 3600
            d_s = abs(s_approx - s) / 1000
            d_time = end - start
            example_results[f"GHA2_approx_{ds:.3f}"] = (d_alpha0, d_alpha1, d_s, d_time)
        except Exception as e:
            print(e)
            example_results[f"GHA2_approx_{ds:.3f}"] = (nan, nan, nan, nan)

    results_panou[f"beta0: {wu.rad2deg(beta0):.3f}, lamb0: {wu.rad2deg(lamb0):.3f}, alpha0: {wu.rad2deg(alpha0_ell):.3f}, s: {s}"] = example_results

print(results_karney)
with open("results_karney.pkl", "wb") as f:
    pickle.dump(results_karney, f)

print(results_panou)
with open("results_panou.pkl", "wb") as f:
    pickle.dump(results_panou, f)