import numpy as np
import pytest

from ekore import harmonics

from . import sx as hsx

# until https://github.com/numba/numba/pull/5660 is confirmed
# we need to deactivate numba prior running


def test_cern_polygamma():
    zs = [1.0, 2.0, 3.0, 0 + 1j, -1 + 1j, -2 + 1j, -1 + 2j, -2 + 2j, -3 + 2j]
    ks = range(5)

    fortran_ref = [
        [
            (-0.5772156649015332 + 0j),
            (0.42278433509846636 + 0j),
            (0.9227843350984672 + 0j),
            (0.09465032062247669 + 2.0766740474685816j),
            (0.5946503206224767 + 2.5766740474685816j),
            (0.9946503206224772 + 2.7766740474685814j),
            (0.9145915153739776 + 2.2208072826422303j),
            (1.1645915153739772 + 2.47080728264223j),
            (1.395360746143208 + 2.624653436488384j),
        ],
        [
            (1.6449340668482264 + 0j),
            (0.6449340668482264 + 0j),
            (0.3949340668482264 + 0j),
            (-0.5369999033772361 - 0.7942335427593189j),
            (-0.5369999033772362 - 0.2942335427593189j),
            (-0.4169999033772362 - 0.13423354275931887j),
            (-0.24506883785905695 - 0.3178255501472297j),
            (-0.24506883785905695 - 0.19282555014722969j),
            (-0.21548303904248892 - 0.12181963298746643j),
        ],
        [
            (-2.404113806319188 + 0j),
            (-0.40411380631918853 + 0j),
            (-0.15411380631918858 + 0j),
            (0.3685529315879351 - 1.233347149654934j),
            (-0.13144706841206477 - 0.7333471496549337j),
            (-0.09944706841206462 - 0.5573471496549336j),
            (0.03902435405364951 - 0.15743252404131272j),
            (-0.02347564594635048 - 0.09493252404131272j),
            (-0.031668636387861625 - 0.053057239562477945j),
        ],
        [
            (6.49393940226683 + 0j),
            (0.49393940226682925 + 0j),
            (0.11893940226682913 + 0j),
            (4.4771255510465044 - 0.31728657866196064j),
            (2.9771255510464925 - 0.3172865786619599j),
            (2.909925551046492 - 0.08688657866195917j),
            (0.12301766661068443 - 0.05523068481179527j),
            (0.02926766661068438 - 0.055230684811795216j),
            (0.004268541930176011 - 0.03002148345329936j),
        ],
        [
            (-24.88626612344089 + 0j),
            (-0.8862661234408784 + 0j),
            (-0.13626612344087824 + 0j),
            (3.2795081690440493 + 21.41938794863803j),
            (0.2795081690440445 + 18.419387948637894j),
            (-0.012331830955960252 + 18.734267948637896j),
            (0.14223316576854003 + 0.10023607930398608j),
            (0.04848316576854002 + 0.006486079303986134j),
            (0.009893695996688708 + 0.014372034600746361j),
        ],
    ]
    for nk, k in enumerate(ks):
        for nz, z in enumerate(zs):
            me = harmonics.polygamma.cern_polygamma(z, k)
            ref = fortran_ref[nk][nz]
            np.testing.assert_almost_equal(me, ref)
    # errors
    with pytest.raises(NotImplementedError):
        _ = harmonics.polygamma.cern_polygamma(1, 5)
    with pytest.raises(ValueError):
        _ = harmonics.polygamma.cern_polygamma(0, 0)


def test_recursive_harmonic_sum():
    n = np.random.rand()
    iterations = 1
    sx_base = hsx(n)
    sx_test = hsx(n + iterations)

    sx_final = []
    for w in range(1, 6):
        sx_final.append(
            harmonics.polygamma.recursive_harmonic_sum(sx_base[w - 1], n, iterations, w)
        )

    np.testing.assert_allclose(sx_final, sx_test)
