"""
This script implements the algorithm of Pitoun and Varescon to test if a number field is p-rational. Reference : Pitoun, F., & Varescon, F. (2015). Computing the torsion of the 𝑝-ramified module of a number field. Mathematics of Computation, 84(291), 371-383.
"""


"""
Compute the invariant factors as in Corollary 4.1 in the reference article.
"""
def FI(K,p,n):
    f=K.defining_polynomial()
    r1,r2=K.signature()
    ab=pari(' K = bnfinit(' + str(f) + ',1);   '+\
           ' bnfcertify(K);               '+\
           ' Kr = bnrinit(K, ' + str(p^n)+ '); '+\
           ' Kr.clgp.cyc ')                       
    # Kr is the Ray class group.
    return ab        
    # return val(ai) and val(bj)
    # where A_{p^{n}}(K) = Z/a1 x ... x Z/a_r2+1 x Z/b1 x ... x Z/bt
    # and val_p(a1) >= ... >=val_p(a_r2+1) >= val_p(b1) >= ...


"""
Test is the number field of f is p-rational. If this number field doesn't verify Leopoldt's conjecture then the programme doesn't terminate.

"""
def is_p_rational(f,p):
    Zx=f.parent()
    K.<a>=NumberField(f)
    r1,r2=K.signature()
    OK=K.ring_of_integers()
    factorization_p=factor(p*OK) # pairs (pi,vi)
    e=max([pivi[1] for pivi in factorization_p]) # second component of (pi,vi)
    s=valuation(e,p)
    n=2+s
    old_ab=FI(K,p,n)
    old_a=FI(K,p,n)[:r2+1]  # first r2+1 components returned by FI
    #                         old_a=[val_p(a1),val_p(a2),...,val_p(a_r2+1)]
    old_b=FI(K,p,n)[r2+1:]  # old_b=[val_p(b1),val_p(b2),...,val_p(bt)]
    n+=1
    found=false
    while not found:
        new_ab=FI(K,p,n)        
        new_a=FI(K,p,n)[:r2+1] # similar to old_a, corresponds to n+1
        new_b=list(FI(K,p,n)[r2+1:]) # similar to old_b, corresponds to n+1
        if new_a ==  [p*ai for ai in old_a] and \
                min(new_a) > p*max(new_b+[1]):
            #  if new_b is empty we replace max(new_b) by 1
            found=true
            if [valuation(bi,p) for bi in new_b] == len(new_b)*[0]: # the elements of new_b are non-negative
                answer=true       # their sum is 0 if they are all zero
            else:
                answer=false
        old_ab=new_ab             # increase n by 1
        old_a=new_a 
        n+=1
    return answer