"""Character-orbit context and standard Sage cross-check -- numberdb.org/T666.
Run: sage -python context_checks.py
The JSON records rational character-orbit factors with multiplicities. These
are cohomological orbit factors, not assertions that their Jacobians are simple.
"""
import json
from math import gcd
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.modular.dirichlet import DirichletGroup
from generate import FermatZeta, character_sums, zeta_polynomial


def standard_sage(N,p):
    """Return the zeta numerator for N >= 3 and split primes p = 1 mod N."""
    N,p = ZZ(N),ZZ(p)
    if N < 3 or not p.is_prime(proof=True) or (p-1) % N:
        raise ValueError('This algorithm requires N >= 3 and a prime p = 1 mod N')
    D=DirichletGroup(p)
    chi=D.gen()**((p-1)//N)
    S=PolynomialRing(D.base_ring(),'T')
    T=S.gen()
    P=S.one()
    for a in range(1,N):
        for b in range(1,N):
            if a+b != N:
                P *= 1+(chi**a).jacobi_sum(chi**b)*T
    return PolynomialRing(ZZ,'T')(P)


def context(N,p):
    K,sums=character_sums(N,p)
    S=PolynomialRing(K,'T'); T=S.gen(); R=PolynomialRing(ZZ,'T')
    remaining=set(sums); orbits=[]
    while remaining:
        a,b=min(remaining)
        orbit=sorted({(h*a % N,h*b % N) for h in range(1,N) if gcd(h,N)==1})
        assert set(orbit) <= remaining
        remaining.difference_update(orbit)
        factor=S.one()
        for pair in orbit: factor*=1+sums[pair]*T
        exact=R([ZZ(c) for c in factor.list()])
        orbits.append({'representative':[a,b],'pairs':[list(pair) for pair in orbit],
                       'd':N//gcd(gcd(N,a),b),'holomorphic_pairs':[list(pair) for pair in orbit if sum(pair)<N],
                       'factor_coefficients':[int(c) for c in exact]})
    product=R.one()
    for orbit in orbits: product*=R(orbit['factor_coefficients'])
    assert product==zeta_polynomial(N,p)
    return {'N':N,'p':p,'orbits':orbits}


if __name__=='__main__':
    rows=[]
    for params in FermatZeta().enumerate():
        N,p=params['N'],params['p']
        assert standard_sage(N,p)==zeta_polynomial(N,p)
        rows.append(context(N,p))
    print('Standard Sage DirichletGroup and all character-orbit products agree on',len(rows),'rows')
    print('CONTEXT_JSON '+json.dumps({'table':'T666','rows':rows,'interpretation':'Full cohomology partitioned into Galois character orbits; equal factors retain separate occurrences; no simplicity claim.'},sort_keys=True))