"""decouple.py -- exact verification of the unit-interval lemma and the conductor-1 decoupling:
sigma_q(x/q) is constant on unit intervals (q>1); <B1^a Psi^b> = <B1^a><Psi^b> exactly;
and the three mixed brackets of the prime-arrival theorem vanish exactly (p=5, mixed Psi).
"""
from fractions import Fraction as F
from math import lcm, gcd

MUv={1:1,2:-1,3:-1,5:-1,6:1,7:-1,10:1,15:1,30:-1,11:-1,13:-1,14:1,21:1,22:1,26:1,33:1,35:1}
def mu(n):
    if n==1: return 1
    r=1; m=n; p=2
    while p*p<=m:
        if m%p==0:
            m//=p
            if m%p==0: return 0
            r=-r
        p+=1
    if m>1: r=-r
    return r

def sigma_val(q, j):
    """value of sigma_q(x/q) on (j, j+1), q>1 (constant there); use midpoint x=j+1/2"""
    x=F(2*j+1,2)
    v=F(0)
    for t in range(1,q+1):
        if q%t==0 and mu(t)!=0:
            y=F(t,1)*x/q   # t*x/q
            v+=F(mu(t),t)*( y - (y.numerator//y.denominator) - F(1,2) )
    return v

# CHECK 0: sigma_q(x/q) constant on unit intervals: compare values at two points inside (j,j+1)
def sigma_at(q, x):
    v=F(0)
    for t in range(1,q+1):
        if q%t==0 and mu(t)!=0:
            y=F(t)*x/q
            v+=F(mu(t),t)*( y - (y.numerator//y.denominator) - F(1,2) )
    return v
ok=True
for q in [2,3,5,6,10,15,30]:
    for j in range(0, 2*q):
        a=sigma_at(q, F(4*j+1,4)); b=sigma_at(q, F(4*j+3,4))
        if a!=b: ok=False; print("NOT CONSTANT", q, j)
print("check0 constant-on-unit-intervals:", ok)

# Build Psi for D: Psi = sum_{1<q<=D} mu(q) m_q(D//q) sigma_q(x/q), piecewise const on (j,j+1), period L=lcm
def mq(q,Y):
    return sum(F(mu(m),m) for m in range(1,int(Y)+1) if gcd(m,q)==1 and mu(m)!=0)

def psi_values(D):
    L=lcm(*range(1,D+1))
    qs=[q for q in range(2,D+1) if mu(q)!=0]
    cs={q: mu(q)*mq(q, D//q) for q in qs}
    vals=[sum(cs[q]*sigma_val(q,j) for q in qs) for j in range(L)]
    return vals, L

# CHECK 1: decoupling <B1^a Psi^b> = <B1^a><Psi^b> for a=1..3, b=1..3, D=6
D=6
vals,L=psi_values(D)
B1m={0:F(1),1:F(0),2:F(1,12),3:F(0),4:F(1,80)}  # int_0^1 (t-1/2)^m dt
for a in range(1,4):
    for b in range(1,4):
        lhs=sum(v**b for v in vals)*B1m[a]/L      # since Psi const per interval: <B1^a Psi^b> = (1/L) sum_j Psi_j^b * int B1^a
        # independent check by fine exact integration: B1^a is polynomial per interval -> integral = Psi_j^b * B1m[a]; same thing. Instead verify via Fourier-free direct Riemann at high resolution:
        pass
# real check: compute <B1^a Psi^b> by exact piecewise integration WITHOUT assuming Psi constant (recompute Psi on subgrid)
def psi_at(D, x):
    s=F(0)
    for q in range(2,D+1):
        if mu(q)!=0:
            s+= mu(q)*mq(q,D//q)*sigma_at(q,x)
    return s
import itertools
L=lcm(*range(1,D+1))
# subdivide each unit interval into K pieces; B1 linear, Psi should be constant; do exact Gauss? use exact integral: on each subinterval Psi const (verified check0) so fine.
for a in [1,2,3]:
    for b in [1,2,3]:
        tot=F(0)
        for j in range(L):
            pv=psi_at(D, F(2*j+1,2))
            # int_j^{j+1} B1^a dx = B1m[a]
            tot+= pv**b * B1m[a]
        lhs=tot/L
        rhs=B1m[a]*sum(psi_at(D,F(2*j+1,2))**b for j in range(L))/L
        print(f"a={a} b={b}  <B1^aPsi^b>={lhs}  <B1^a><Psi^b>={rhs}  equal={lhs==rhs}")

# CHECK 2: prime-arrival brackets for p=5 against Psi built from conductors {2,3} (weights 1), period lcm(5,6)=30
p=5
Lp=30
sig5=[sigma_val(5,j%5) for j in range(Lp)]
Ps=[sigma_val(2,j%2)+sigma_val(3,j%3) for j in range(Lp)]  # arbitrary nonzero weights 1,1
m1=lambda f: sum(f)/Lp
br1=m1([s*w**3 for s,w in zip(sig5,Ps)])
br2=m1([s**2*w**2 for s,w in zip(sig5,Ps)]) - m1([s**2 for s in sig5])*m1([w**2 for w in Ps])
br3=m1([s**3*w for s,w in zip(sig5,Ps)])
print("CHECK2 p=5:", br1, br2, br3)

# CHECK 3: same with Psi including composite conductor 6: Psi=2*sig2+3*sig3-sig6 (random weights)
Ps=[2*sigma_val(2,j%2)+3*sigma_val(3,j%3)-sigma_val(6,j%6) for j in range(Lp)]
br1=m1([s*w**3 for s,w in zip(sig5,Ps)])
br2=m1([s**2*w**2 for s,w in zip(sig5,Ps)]) - m1([s**2 for s in sig5])*m1([w**2 for w in Ps])
br3=m1([s**3*w for s,w in zip(sig5,Ps)])
print("CHECK3 p=5, Psi w/ conductor 6:", br1, br2, br3)