#!/usr/bin/env python3
"""
Fast high-precision Python verification code for the constants in
"Thinned Wallis-type prime products in residue classes modulo 2^m".

It evaluates the character-sum formula for log K(2^m,S).  The expensive
Dirichlet L-values are computed for all characters at once by a finite
Fourier transform on (Z/2^m Z)^* = {+-1} x <5>.

Dependencies: Python 3 and mpmath.
"""
from __future__ import annotations

import argparse
import math
from typing import Dict, List, Sequence, Tuple
import mpmath as mp

Character = Tuple[int, int]  # e,r with chi(-1)=(-1)^e, chi(5)=exp(2*pi*i*r/N)


def mobius_sieve(n: int) -> List[int]:
    mu = [0] * (n + 1)
    if n >= 1:
        mu[1] = 1
    primes: List[int] = []
    is_comp = [False] * (n + 1)
    for i in range(2, n + 1):
        if not is_comp[i]:
            primes.append(i)
            mu[i] = -1
        for p in primes:
            ip = i * p
            if ip > n:
                break
            is_comp[ip] = True
            if i % p == 0:
                mu[ip] = 0
                break
            mu[ip] = -mu[i]
    return mu


def divisor_table(nmax: int) -> List[List[int]]:
    divs = [[] for _ in range(nmax + 1)]
    for d in range(1, nmax + 1):
        for n in range(d, nmax + 1, d):
            divs[n].append(d)
    return divs


class PowerTwoCharacters:
    def __init__(self, m: int):
        if m < 2:
            raise ValueError("Require m>=2.")
        self.m = m
        self.q = 1 << m
        self.phi = 1 << (m - 1)
        if m == 2:
            self.N = 1
            self.a0 = [1]
            self.a1 = [3]
            self.logmap = {1: (0, 0), 3: (1, 0)}
        else:
            self.N = 1 << (m - 2)
            self.a0 = []
            self.a1 = []
            self.logmap = {}
            cur = 1
            for t in range(self.N):
                self.a0.append(cur)
                self.a1.append((-cur) % self.q)
                self.logmap[cur] = (0, t)
                self.logmap[(-cur) % self.q] = (1, t)
                cur = (cur * 5) % self.q
        self.characters: List[Character] = [(e, r) for e in (0, 1) for r in range(self.N)]

    @staticmethod
    def chi4(a: int) -> int:
        return 1 if a % 4 == 1 else -1

    @staticmethod
    def principal() -> Character:
        return (0, 0)

    @staticmethod
    def is_principal(ch: Character) -> bool:
        return ch == (0, 0)

    def power(self, ch: Character, j: int) -> Character:
        e, r = ch
        return ((e * j) & 1, (r * j) % self.N if self.N > 1 else 0)

    def value(self, ch: Character, a: int) -> mp.mpc:
        a %= self.q
        if a % 2 == 0:
            return mp.mpc(0)
        e, r = ch
        sign_part, t = self.logmap[a]
        val = -1 if (e and sign_part) else 1
        if self.N > 1:
            val *= mp.e ** (2j * mp.pi * r * t / self.N)
        return mp.mpc(val)


def fft_pos(vec):
    """Radix-2 FFT with positive exponent: F[k]=sum_j vec[j]*exp(2*pi*i*k*j/N)."""
    n = len(vec)
    if n == 1:
        return [mp.mpc(vec[0])]
    even = fft_pos(vec[0::2])
    odd = fft_pos(vec[1::2])
    out = [mp.mpc(0)] * n
    for k in range(n // 2):
        tw = mp.e ** (2j * mp.pi * k / n)
        t = tw * odd[k]
        out[k] = even[k] + t
        out[k + n // 2] = even[k] - t
    return out


class LValueTable:
    def __init__(self, group: PowerTwoCharacters, max_s: int):
        self.G = group
        self.max_s = max_s
        self.table: Dict[Tuple[Character, int], mp.mpc] = {}
        self._build()

    def _build(self) -> None:
        G = self.G
        qmp = mp.mpf(G.q)
        for s in range(1, self.max_s + 1):
            F0: List[mp.mpc] = []
            F1: List[mp.mpc] = []
            if s == 1:
                for a in G.a0:
                    F0.append(-mp.digamma(mp.mpf(a) / G.q) / G.q)
                for a in G.a1:
                    F1.append(-mp.digamma(mp.mpf(a) / G.q) / G.q)
            else:
                qpow = qmp ** s
                for a in G.a0:
                    F0.append(mp.zeta(s, mp.mpf(a) / G.q) / qpow)
                for a in G.a1:
                    F1.append(mp.zeta(s, mp.mpf(a) / G.q) / qpow)

            plus_transform = fft_pos([F0[t] + F1[t] for t in range(G.N)])
            minus_transform = fft_pos([F0[t] - F1[t] for t in range(G.N)])
            for r in range(G.N):
                self.table[((0, r), s)] = plus_transform[r]
                self.table[((1, r), s)] = minus_transform[r]

            # Replace the principal character by the exact Euler-product expression for s>1.
            # For s=1 the pole is intentionally left undefined; it is never used.
            if s > 1:
                self.table[(G.principal(), s)] = mp.zeta(s) * (1 - mp.power(2, -s))

    def L(self, ch: Character, s: int) -> mp.mpc:
        if s == 1 and self.G.is_principal(ch):
            raise ZeroDivisionError("principal character has a pole at s=1")
        return self.table[(ch, int(s))]


def logK(m: int, S: Sequence[int], nmax: int = 40, dps: int = 80) -> mp.mpc:
    mp.mp.dps = dps
    G = PowerTwoCharacters(m)
    S0 = [a % G.q for a in S]
    for a in S0:
        if math.gcd(a, G.q) != 1:
            raise ValueError(f"{a} is not a reduced residue class modulo {G.q}.")

    Ltab = LValueTable(G, 2 * nmax)
    mu = mobius_sieve(nmax)
    divs = divisor_table(nmax)
    balance = sum(G.chi4(a) for a in S0)

    total = mp.mpc(0)
    if balance != 0:
        principal_part = 2 * (mp.log(2) - mp.euler) + mp.log(Ltab.L(G.principal(), 2))
        total += mp.power(2, 1 - m) * balance * principal_part

    for ch in G.characters:
        if G.is_principal(ch):
            continue
        coeff = mp.mpc(0)
        for a in S0:
            coeff += G.chi4(a) * G.value(ch, a)
        if abs(coeff) < mp.mpf(10) ** (-(dps - 10)):
            continue

        inner = mp.mpc(0)
        for n in range(1, nmax + 1):
            subtotal = mp.mpc(0)
            for j in divs[n]:
                muj = mu[j]
                if muj == 0:
                    continue
                chj = G.power(ch, j)
                subtotal += muj * (mp.log(Ltab.L(chj, 2 * n)) - 2 * mp.log(Ltab.L(chj, n)))
            inner += subtotal / n
        total += mp.power(2, 1 - m) * coeff * inner
    return total


def sieve_primes_upto(x: int) -> List[int]:
    if x < 2:
        return []
    sieve = bytearray(b"\x01") * (x + 1)
    sieve[0:2] = b"\x00\x00"
    r = math.isqrt(x)
    for p in range(2, r + 1):
        if sieve[p]:
            sieve[p * p: x + 1: p] = b"\x00" * (((x - p * p) // p) + 1)
    return [i for i in range(2, x + 1) if sieve[i]]


def direct_log_product(m: int, S: Sequence[int], x: int, dps: int = 50) -> mp.mpf:
    mp.mp.dps = dps
    q = 1 << m
    Sset = {a % q for a in S}
    total = mp.mpf(0)
    for p in sieve_primes_upto(x):
        if p != 2 and p % q in Sset:
            c = 1 if p % 4 == 1 else -1
            total += mp.log(mp.mpf(p - c) / mp.mpf(p + c))
    return total


def print_constant(m: int, S: Sequence[int], nmax: int, dps: int) -> None:
    v = logK(m, S, nmax=nmax, dps=dps)
    print(f"m={m}, q={1<<m}, S={list(S)}, nmax={nmax}, dps={dps}")
    print("logK =", mp.nstr(mp.re(v), max(20, dps - 10)))
    print("Im(logK) =", mp.nstr(mp.im(v), 8))
    print("K =", mp.nstr(mp.e ** mp.re(v), max(20, dps - 10)))


def print_stability(m: int, S: Sequence[int], nvals: Sequence[int], dps: int) -> None:
    print(f"Stability for m={m}, q={1<<m}, S={list(S)}, dps={dps}")
    for n in nvals:
        v = logK(m, S, nmax=n, dps=dps)
        print(f"  nmax={n:3d}: logK={mp.nstr(mp.re(v), 45)}  K={mp.nstr(mp.e**mp.re(v), 35)}  Im={mp.nstr(mp.im(v), 5)}")


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--m", type=int, default=4)
    p.add_argument("--S", type=int, nargs="+", default=[1, 15])
    p.add_argument("--nmax", type=int, default=40)
    p.add_argument("--dps", type=int, default=80)
    p.add_argument("--stability", action="store_true")
    p.add_argument("--nvals", type=int, nargs="+", default=[16, 24, 32, 40])
    p.add_argument("--direct", action="store_true")
    p.add_argument("--x", type=int, default=10_000_000)
    args = p.parse_args()

    if args.direct:
        lp = direct_log_product(args.m, args.S, args.x, args.dps)
        balance = sum(1 if a % 4 == 1 else -1 for a in args.S)
        exponent = -2 * balance / (1 << (args.m - 1))
        P = mp.e ** lp
        norm = P / (mp.log(args.x) ** exponent)
        print(f"m={args.m}, q={1<<args.m}, S={args.S}, x={args.x}")
        print("P_S(x) =", mp.nstr(P, max(20, args.dps - 10)))
        print("normalized =", mp.nstr(norm, max(20, args.dps - 10)))
    elif args.stability:
        print_stability(args.m, args.S, args.nvals, args.dps)
    else:
        print_constant(args.m, args.S, args.nmax, args.dps)


if __name__ == "__main__":
    main()