"""wcollapse.py -- collapsed cotangent formula for W_{r,s}(xi).

Theorem (to verify): for squarefree r != s >1, g=(r,s), r1=r/g, s1=s/g,
xi with N = rs*xi in Z, N != 0:

  W(xi) = (pi/N) * [ Sum_{a in Au} cot(pi a / (g r1))  +  Sum_{b in Bv} cot(pi b/(g s1)) ]

where Au = { a mod g*r1 :  a*s === N (mod r)  [one class mod r1],
                           (a, g*r1) = 1,
                           v(a) := (N - a s)/r  is a unit mod every p | g }
and Bv the mirror (swap roles of (u,r) and (v,s)).

Derivation: the solution classes of u live mod M=rs but all defining
conditions only depend on a mod g*r1; each class mod g*r1 is a full AP of
modulus g*r1 inside Z/M, and sum_{j mod n} cot(pi(x+j)/n) = n cot(pi x)
collapses it. The prefactors: r * (pi/M) * (g s1) = pi, same for s.
"""
import numpy as np
from math import gcd, pi

def W_collapsed(r, s, xi_num, xi_den):
    g = gcd(r, s); r1, s1 = r // g, s // g
    M = r * s
    N = M * xi_num // xi_den
    assert M * xi_num % xi_den == 0 and N != 0
    if gcd(N, r1 * s1) != 1:
        return 0.0   # a prime of r1 (resp. s1) dividing N kills both sides
    tot = 0.0
    # u-side: a mod g*r1
    mu_ = g * r1
    for a in range(mu_):
        if gcd(a, mu_) != 1: continue          # (a, r)=1 as class mod g*r1
        if (a * s - N) % r: continue           # a*s === N (mod r); note r | (as-N) iff class ok
        v = (N - a * s) // r                   # integer for any lift? need as===N mod r: use lift
        # v mod p for p|g depends on a mod g*r1 only; compute with this lift
        if gcd(v % s if v % s else s, g) != 1 and g > 1:
            pass
        # cleaner: check gcd(v, g) using the lift value mod g
        if g > 1 and gcd(v % g if v % g else g, g) != 1: continue
        tot += 1.0 / np.tan(pi * a / mu_)
    # v-side: b mod g*s1
    mv = g * s1
    for b in range(mv):
        if gcd(b, mv) != 1: continue
        if (b * r - N) % s: continue
        u = (N - b * r) // s
        if g > 1 and gcd(u % g if u % g else g, g) != 1: continue
        tot2 = 1.0 / np.tan(pi * b / mv)
        tot += tot2
    return (pi / N) * tot

if __name__ == "__main__":
    from wlaw import W_cot
    print("== validation: collapsed vs full cot formula ==")
    cases = [(10,15,1,6),(10,15,5,6),(6,15,1,10),(21,35,1,15),(21,35,2,15),
             (10,15,1,30),(14,35,3,10),(22,33,1,6),(2*97,3*97,1,6),(115,145,2,667),
             (30,105,1,14),(26,39,1,6),(58,62,1,899),(101,103,1,10403)]
    for r,s,kn,kd in cases:
        if (r*s*kn) % kd: print("skip", r,s,kn,kd); continue
        wc = W_cot(r, s, kn, kd)
        wf = W_collapsed(r, s, kn, kd)
        print(f"r={r:4d} s={s:4d} xi={kn}/{kd:6d}: full={wc:+.10f} collapsed={wf:+.10f} diff={abs(wc-wf):.2e}")