import adani as ad
import oldadani as old
import numpy as np

def test_nll():
    for nll in [True, False]:
        he = ad.HighEnergyCoefficientFunction(3, "2", "g", NLL=nll)
        hehs = ad.HighEnergyHighScaleCoefficientFunction(3, "2", "g", NLL=nll)
        pt = ad.PowerTermsCoefficientFunction(3, "2", "g", NLL=nll)
        assert he.GetNLL() == nll
        assert hehs.GetNLL() == nll
        assert pt.GetNLL() == nll

def test_highenergy_as2_oldadani():
    for kind in ['2', 'L']:
        for channel in ['g', 'q']:
            he = ad.HighEnergyCoefficientFunction(2, kind, channel)
            for x in np.geomspace(1e-5, 1., 100, endpoint=False):
                for m2Q2 in np.geomspace(1e-4, 1e-2, 10):
                    for m2mu2 in np.geomspace(1e-4, 1e-2, 10):
                        res1 = he.fx(x, m2Q2, m2mu2, 0)
                        if (kind == '2' and channel == 'g'):
                            res2 = old.C2_g2_highenergy(x, m2Q2, m2mu2)
                        if (kind == '2' and channel == 'q'):
                            res2 = old.C2_ps2_highenergy(x, m2Q2, m2mu2)
                        if (kind == 'L' and channel == 'g'):
                            res2 = old.CL_g2_highenergy(x, m2Q2, m2mu2)
                        if (kind == 'L' and channel == 'q'):
                            res2 = old.CL_ps2_highenergy(x, m2Q2, m2mu2)
                        np.testing.assert_allclose(res1, res2, rtol=1e-7)

def test_highenergy_as3_oldadani():
    for kind in ['2', 'L']:
        for channel in ['g', 'q']:
            he = ad.HighEnergyCoefficientFunction(3, kind, channel)
            for x in np.geomspace(1e-5, 1., 100, endpoint=False):
                for m2Q2 in np.geomspace(1e-4, 1e-2, 10):
                    for m2mu2 in np.geomspace(1e-4, 1e-2, 10):
                        for nf in range(6 + 1):
                            res1 = he.fx(x, m2Q2, m2mu2, nf)
                            if (kind == '2' and channel == 'g'):
                                res2 = old.C2_g3_highenergy(x, m2Q2, m2mu2, nf, 0)
                            if (kind == '2' and channel == 'q'):
                                res2 = old.C2_ps3_highenergy(x, m2Q2, m2mu2, nf, 0)
                            if (kind == 'L' and channel == 'g'):
                                res2 = old.CL_g3_highenergy(x, m2Q2, m2mu2, nf, 0)
                            if (kind == 'L' and channel == 'q'):
                                res2 = old.CL_ps3_highenergy(x, m2Q2, m2mu2, nf, 0)
                            np.testing.assert_allclose(res1, res2, rtol=1e-7)

def test_highenergyhighscale_as2_oldadani():
    for kind in ['2', 'L']:
        for channel in ['g', 'q']:
            he = ad.HighEnergyHighScaleCoefficientFunction(2, kind, channel)
            for x in np.geomspace(1e-5, 1., 100, endpoint=False):
                for m2Q2 in np.geomspace(1e-4, 1e-2, 10):
                    for m2mu2 in np.geomspace(1e-4, 1e-2, 10):
                        res1 = he.fx(x, m2Q2, m2mu2, 0)
                        if (kind == '2' and channel == 'g'):
                            res2 = old.C2_g2_highenergy_highscale(x, m2Q2, m2mu2)
                        if (kind == '2' and channel == 'q'):
                            res2 = old.C2_ps2_highenergy_highscale(x, m2Q2, m2mu2)
                        if (kind == 'L' and channel == 'g'):
                            res2 = old.CL_g2_highenergy_highscale(x, m2Q2, m2mu2)
                        if (kind == 'L' and channel == 'q'):
                            res2 = old.CL_ps2_highenergy_highscale(x, m2Q2, m2mu2)
                        np.testing.assert_allclose(res1, res2, rtol=1e-7)

def test_highenergyhighscale_as3_oldadani():
    for kind in ['2', 'L']:
        for channel in ['g', 'q']:
            he = ad.HighEnergyHighScaleCoefficientFunction(3, kind, channel)
            for x in np.geomspace(1e-5, 1., 100, endpoint=False):
                for m2Q2 in np.geomspace(1e-4, 1e-2, 10):
                    for m2mu2 in np.geomspace(1e-4, 1e-2, 10):
                        for nf in range(6 + 1):
                            res1 = he.fx(x, m2Q2, m2mu2, nf)
                            if (kind == '2' and channel == 'g'):
                                res2 = old.C2_g3_highenergy_highscale(x, m2Q2, m2mu2, nf, 0)
                            if (kind == '2' and channel == 'q'):
                                res2 = old.C2_ps3_highenergy_highscale(x, m2Q2, m2mu2, nf, 0)
                            if (kind == 'L' and channel == 'g'):
                                res2 = old.CL_g3_highenergy_highscale(x, m2Q2, m2mu2, nf, 0)
                            if (kind == 'L' and channel == 'q'):
                                res2 = old.CL_ps3_highenergy_highscale(x, m2Q2, m2mu2, nf, 0)
                            np.testing.assert_allclose(res1, res2, rtol=1e-7)

def test_powerterms_as2_oldadani():
    for kind in ['2', 'L']:
        for channel in ['g', 'q']:
            he = ad.PowerTermsCoefficientFunction(2, kind, channel)
            for x in np.geomspace(1e-5, 1., 100, endpoint=False):
                for m2Q2 in np.geomspace(1e-4, 1e-2, 10):
                    for m2mu2 in np.geomspace(1e-4, 1e-2, 10):
                        res1 = he.fx(x, m2Q2, m2mu2, 0)
                        if (kind == '2' and channel == 'g'):
                            res2 = old.C2_g2_power_terms(x, m2Q2, m2mu2)
                        if (kind == '2' and channel == 'q'):
                            res2 = old.C2_ps2_power_terms(x, m2Q2, m2mu2)
                        if (kind == 'L' and channel == 'g'):
                            res2 = old.CL_g2_power_terms(x, m2Q2, m2mu2)
                        if (kind == 'L' and channel == 'q'):
                            res2 = old.CL_ps2_power_terms(x, m2Q2, m2mu2)
                        np.testing.assert_allclose(res1, res2, rtol=1e-7)

def test_powerterms_as3_oldadani():
    for kind in ['2', 'L']:
        for channel in ['g', 'q']:
            he = ad.PowerTermsCoefficientFunction(3, kind, channel)
            for x in np.geomspace(1e-5, 1., 100, endpoint=False):
                for m2Q2 in np.geomspace(1e-4, 1e-2, 10):
                    for m2mu2 in np.geomspace(1e-4, 1e-2, 10):
                        for nf in range(6 + 1):
                            res1 = he.fx(x, m2Q2, m2mu2, nf)
                            if (kind == '2' and channel == 'g'):
                                res2 = old.C2_g3_power_terms(x, m2Q2, m2mu2, nf, 0)
                            if (kind == '2' and channel == 'q'):
                                res2 = old.C2_ps3_power_terms(x, m2Q2, m2mu2, nf, 0)
                            if (kind == 'L' and channel == 'g'):
                                res2 = old.CL_g3_power_terms(x, m2Q2, m2mu2, nf, 0)
                            if (kind == 'L' and channel == 'q'):
                                res2 = old.CL_ps3_power_terms(x, m2Q2, m2mu2, nf, 0)
                            np.testing.assert_allclose(res1, res2, rtol=1e-7)

def test_powerterms():
    for order in range(1, 3 + 1):
        for kind in ['2', 'L']:
            for channel in ['g', 'q']:
                if channel == 'q' and order == 1:
                    continue
                for NLL in [True, False]:
                    he = ad.HighEnergyCoefficientFunction(order, kind, channel, NLL)
                    hehs = ad.HighEnergyHighScaleCoefficientFunction(order, kind, channel, NLL)
                    pt = ad.PowerTermsCoefficientFunction(order, kind, channel, NLL)
                    for nf in range(1, 6 + 1):
                        for m2mu2 in np.geomspace(1e-4, 1e2, 10):
                            for m2Q2 in np.geomspace(1e-4, 1e2, 10):
                                for x in np.geomspace(1e-5, 1., 10, endpoint=False):
                                    res1 = pt.fxBand(x, m2Q2, m2mu2, nf)

                                    res2_c = he.fxBand(x, m2Q2, m2mu2, nf).GetCentral() - hehs.fxBand(x, m2Q2, m2mu2, nf).GetCentral()
                                    tmp1 = he.fxBand(x, m2Q2, m2mu2, nf).GetHigher() - hehs.fxBand(x, m2Q2, m2mu2, nf).GetHigher()
                                    tmp2 = he.fxBand(x, m2Q2, m2mu2, nf).GetLower() - hehs.fxBand(x, m2Q2, m2mu2, nf).GetLower()
                                    if tmp1 > tmp2:
                                        res2_h = tmp1
                                        res2_l = tmp2
                                    else:
                                        res2_h = tmp2
                                        res2_l = tmp1

                                    np.testing.assert_allclose(res1.GetCentral(), res2_c, rtol=1e-7)
                                    np.testing.assert_allclose(res1.GetHigher(), res2_h, rtol=1e-7)
                                    np.testing.assert_allclose(res1.GetLower(), res2_l, rtol=1e-7)
                    del he
                    del pt
                    del hehs
