import math
import time
import mpmath as mp
from scipy.stats import irwinhall

mp.mp.dps = 60
mp.iv.dps = 60
IVZERO = mp.iv.mpf([0, 0])
IVONE = mp.iv.mpf([1, 1])
IVTWO = mp.iv.mpf([2, 2])
IVFOUR = mp.iv.mpf([4, 4])


def same_cell(a, b):
    return math.floor(a) == math.floor(b)


def g_float(d, u):
    return irwinhall.pdf(u, d) - 2.0 * irwinhall.pdf(2.0 * u, d)


def isolate_root_candidates(d):
    a = d / 4.0
    b = d / 2.0
    hist = []
    for _ in range(80):
        hist.append((a, b))
        m = 0.5 * (a + b)
        fm = g_float(d, m)
        if fm > 0.0:
            b = m
        else:
            a = m
    hist.append((a, b))
    cands = []
    for aa, bb in hist:
        if (
            same_cell(aa, bb)
            and same_cell(2 * aa, 2 * bb)
            and same_cell(aa - 1, bb - 1)
            and same_cell(2 * aa - 1, 2 * bb - 1)
        ):
            cands.append((mp.mpf(aa), mp.mpf(bb)))
    if not cands:
        raise RuntimeError(("could not isolate root into fixed cells", d))
    cands.sort(key=lambda ab: ab[1] - ab[0])
    return cands


def build_rows(d, t, max_j):
    q = [IVONE]
    prev = None
    for n in range(2, d + 1):
        m = min(max_j, n - 1)
        inv = mp.iv.mpf([mp.mpf(1) / (n - 1), mp.mpf(1) / (n - 1)])
        new = [IVZERO] * (m + 1)
        for r in range(m + 1):
            val = IVZERO
            if r < len(q):
                val += (mp.iv.mpf([r, r]) + t) * q[r] * inv
            if r - 1 >= 0:
                val += (mp.iv.mpf([n - r, n - r]) - t) * q[r - 1] * inv
            new[r] = val
        prev, q = q, new
    return prev, q


def rigorous_pack(d, a, b):
    ju = int(mp.floor(a))
    jv = int(mp.floor(2 * a))
    tu = mp.iv.mpf([a - ju, b - ju])
    tv = mp.iv.mpf([2 * a - jv, 2 * b - jv])
    row_dm1_u, row_d_u = build_rows(d, tu, ju)
    row_dm1_v, row_d_v = build_rows(d, tv, jv)
    g = row_d_u[ju] - IVTWO * row_d_v[jv]
    A = row_dm1_v[jv - 1] / row_dm1_u[ju - 1]
    B = row_dm1_u[ju] / (IVTWO * row_dm1_u[ju - 1] + IVFOUR * row_dm1_v[jv])
    return g, A, B


def bounds(ivx):
    return mp.mpf(ivx.a), mp.mpf(ivx.b)


def verify_d(d):
    for a, b in isolate_root_candidates(d):
        gN, A, B = rigorous_pack(d, a, b)
        glo, ghi = bounds(gN)
        # check the interval contains the root u_d^*
        if glo < 0 < ghi:
            Alo, Ahi = bounds(A)
            Blo, Bhi = bounds(B)
            if Alo < 2.01 or Blo < 1.01:
                raise RuntimeError(("inequality failure", d, (a, b), (Alo, Ahi), (Blo, Bhi)))
            return (a, b), (Alo, Ahi), (Blo, Bhi), (glo, ghi)
    raise RuntimeError(("no interval-Newton enclosure found", d))


def main(start=5, stop=200):
    t0 = time.time()
    for d in range(start, stop + 1):
    #for d in [4,5,6,7,8,9,10,20,50,100,150,200]:
        N, A_bounds, B_bounds, _ = verify_d(d)
        print(f"d={d:3d},  A in [{A_bounds[0]}, {A_bounds[1]}],  B in [{B_bounds[0]}, {B_bounds[1]}],  ud in [{N[0]}, {N[1]}]")
        #print(f"    {d:3d} & {float(N[0]):.6f} & {float(A_bounds[0]):.6f} & {float(B_bounds[0]):.6f}\\\\")
    print(f"Elapsed seconds: {time.time()-t0}")


if __name__ == '__main__':
    main(start=4, stop=200)