from mpmath import mp, polyroots
import time

NRUN = 5

def test_polyroots_legendre():
    n = 64
    coeffs = [11975573020964041433067793888190275875, 0,
        -190100434726484311252477736051902332000, 0,
        1437919688271127330313741595496589239248, 0,
        -6897338342113537600691931230430793911840, 0,
        23556405536185284408974715545252277554280, 0,
        -60969520211303089058522793175947071316960, 0,
        124284021969194758465450309166353645376880, 0,
        -204721258548015217049921875719981284186016, 0,
        277415422258095841688223780704620656114900, 0,
        -313237834141273382807123548182995095192800, 0,
        297432255354328395601259515935229287637200, 0,
        -239057700565161140389797367947941296605600, 0,
        163356095386193445933028201431093219347160, 0,
        -95158890516229191805647495979277603503200, 0,
        47310254620162038075933656063247634556400, 0,
        -20071017111583894941305187420771723751200, 0,
        7255051932731034189479516844750603752850, 0,
        -2228176940331017311443863996901733412640, 0,
        579006552594977616773047095969088431600, 0,
        -126584428502545713788439446082310831200, 0,
        23112325428835593809686977515028663000, 0,
        -3491517141958743235617737161547844000, 0,
        431305058712550634988073414073557200, 0,
        -42927166660756742088912492757452000, 0,
        3378527005707706553294038781836500, 0,
        -205277590220215081719131470288800, 0,
        9330799555464321896324157740400, 0,
        -304114948474392713657972548576, 0,
        6695289961520387531608984680, 0,
        -91048139350447232095702560, 0,
        659769125727878493447120, 0,
        -1905929106580294155360, 0,
        916312070471295267]

    with mp.workdps(3):
        try:
            roots = polyroots(coeffs, maxsteps=5, cleanup=True, error=False,
                        extraprec=n*10)
            raise AssertionError("polyroots() didn't raise NoConvergence")
        except (mp.NoConvergence):
            pass

        roots = polyroots(coeffs, maxsteps=50, cleanup=True, error=False,
                    extraprec=n*10)
        roots = [str(r) for r in roots]
        assert roots == \
            ['-0.999', '-0.996', '-0.991', '-0.983', '-0.973', '-0.961',
            '-0.946', '-0.93', '-0.911', '-0.889', '-0.866', '-0.841',
            '-0.813', '-0.784', '-0.753', '-0.72', '-0.685', '-0.649',
            '-0.611', '-0.572', '-0.531', '-0.489', '-0.446', '-0.402',
            '-0.357', '-0.311', '-0.265', '-0.217', '-0.17', '-0.121',
            '-0.073', '-0.0243', '0.0243', '0.073', '0.121', '0.17', '0.217',
            '0.265', '0.311', '0.357', '0.402', '0.446', '0.489', '0.531',
            '0.572', '0.611', '0.649', '0.685', '0.72', '0.753', '0.784',
            '0.813', '0.841', '0.866', '0.889', '0.911', '0.93', '0.946',
            '0.961', '0.973', '0.983', '0.991', '0.996', '0.999']

times = []
for run in range(1, NRUN+1):
    start = time.perf_counter()
    test_polyroots_legendre()
    dt = time.perf_counter() - start
    print("Run #%s/%s: %.2f sec" % (run, NRUN, dt))
    times.append(dt)

print()
print("test_polyroots_legendre (min of %s runs): %.2f sec"
      % (NRUN, min(times)))