from sympy import *
from sympy.abc import q,z
from math import log
import matplotlib.pyplot as plt
init_printing()

def truncate_matrix(a,order):
    size = a.shape[0]
    b = zeros(size)
    for i in range(0,size):
        for j in range(0,size):
            b[i,j] = (a[i,j]+O(q**(order+1))).removeO()
    return b

def p_adic_valuation_int(i):
    if i == 0:
        return float('inf')
    v = 0
    while i % p == 0:
        i = i // p
        v += 1
    return v

def p_adic_valuation(r):
    if r == 0:
        return float('inf')
    a = numer(r)
    b = denom(r)
    return p_adic_valuation_int(a)-p_adic_valuation_int(b)

def matrix_p_adic_valuation(a):
    v = float('inf')
    size0 = a.shape[0]
    size1 = a.shape[1]
    for i in range(0,size0):
        for j in range(0,size1):
            if p_adic_valuation(a[i,j]) < v:
                v = p_adic_valuation(a[i,j])
    return v

def list_of_matrix_p_adic_valuation(l):
    print("computing list of p-adic valuations of",len(l),"terms", end="")
    v = []
    c = 0
    for a in l:
        c += 1
        if c % 10 == 0:
            print(".", end="")
        v.append(matrix_p_adic_valuation(a))
    print()
    return v

def truncate(f,o):
    return (f+O(z**(o+1))).removeO()

def coefficients(f,o):
    l = []
    for i in range(0,o+1):
        l.append(f.coeff(z,i))
    return l

def dwork_exp(o):
    f = 1
    for i in range(1,o+1):
        f+=(z+z**p*Rational(1,p))**i*Rational(1,factorial(i))
    return truncate(expand(f),o)

def dwork_exp_coeff(o):
    """
    computes the Taylor coefficients of exp(x+x^p/p),
    following the recursion idea from Rodriguez-Villegas Experimental Number Theory
    """
    l = [Rational(1,1)]
    for i in range(1,o+1):
        x = l[i-1]
        if i-p >= 0:
            x += l[i-p]
        l.append(x/i)
    return l

def gamma_derivatives(i,val):
    """
    computes the power series of Gamma(x) following
    Rodriguez Villegas p. 156, or Cohen number theory ch. 11
    exactness is guaranteed up to and including p-adic valuation val
    """
    m = 1
    while (m < i*p/((p-1)*log(p)) + 1) or (i/(p-1)+m*(p-1)/p-log(i)/log(p) - i*log(m-1)/log(p) - (2*p-1)/p <= val):
        m += 1
    print("computing order",i,"Gamma derivative up to p-adic valuation",val, end="")
    print(", using",m,"terms in the Mahler expansion", end="")
    b = dwork_exp_coeff(p*m)
    d = 0
    zprod = 1
    for k in range(1,m+1):
        if k % 10 == 0:
            print(".", end="")
        zprod = expand(zprod*(z-k+1))
        new = (-p)**k*b[p*k]*zprod.coeff(z,i)*factorial(i)*Rational(1,-p)**i
        if p_adic_valuation(new) <= val:
            d += new
    print()
    return d

def gamma_derivative(val):
    """
    previous idea specialized to only the first derivative
    """
    m = 1
    while (m < p/((p-1)*log(p)) + 1) or (1/(p-1)+m*(p-1)/p - log(m-1)/log(p) - (2*p-1)/p <= val):
        m += 1
    print("computing first Gamma derivative up to p-adic valuation",val, end="")
    print(", using",m,"terms in the Mahler expansion", end="")
    b = dwork_exp_coeff(p*m)
    d = 0
    for k in range(1,m+1):
        if k % 10 == 0:
            print(".", end="")
        new = p**(k-1)*factorial(k-1)*b[k*p]
        if p_adic_valuation(new) <= val:
            d += new
    print()
    return d

def matrix_to_vector(a):
    """
    This converts a matrix into a column vector, concatenating the rows
    """
    size0 = a.shape[0]
    size1 = a.shape[1]
    v = zeros(size0*size1,1)
    for i in range(0,size0):
        for j in range(0,size1):
            v[i+size0*j,0] = a[i,j]
    return v

def vector_to_matrix(v):
    """
    This converts a column vector of size N^2 to a square matrix of size N
    """
    size = int(sqrt(v.shape[0]+0.01))
    a = zeros(size,size)
    for i in range(0,size):
        for j in range(0,size):
            a[i,j] = v[i+size*j,0]
    return a

def rhs_matrix(givena, givenphi, m):
    """
    This determines the right hand side of the equation
    \partial_q\Phi + A Phi - q^{p-1} Phi A(-q^p/p)  = 0
    At a given order q^{n-1}, n>=1. The givenphi starts with order 0,
    the givena with order -1 (which is not used here), causing an offset in the index
    Each givena and givenphi is a list of matrices with rational coefficients
    The outcome is a single matrix
    """
    size = givena[0].shape[0]
    rhs = zeros(size,size)
    for i in range(0,m):
        if m-i < len(givena):
            rhs -= givena[m-i]*givenphi[i]
        if (m-i) % p == 0:
            k = (m-p-i)//p
            if k < len(givena)-1:
                rhs -= givenphi[i]*givena[k+1]*((-Rational(p,1))**(-k))
    return rhs

def lhs_sample_matrix(givena,phi,m):
    """
    Here, we start with a proposal for phi_m (a matrix), and
    compute the left hand of the equation as a matrix
    """
    b = phi*m
    b += givena[0]*phi - p*phi*givena[0]
    return b

def lhs_matrix(givena,n):
    """
    This tries out all choices for phi_m, and returns the
    left hand side as a matrix of size N^2 \times N^2
    """
    size = givena[0].shape[0]
    lhs = zeros(size**2)
    for i in range(0,size):
        for j in range(0,size):
            phi = zeros(size,size)
            phi[j,i] = 1
            l = matrix_to_vector(lhs_sample_matrix(givena,phi,n))
            for k in range(0,len(l)):
                lhs[k,i*size+j] = l[k,0]
    return lhs

def solve_one_step(givena,givenphi,n):
    rhs_v = matrix_to_vector(rhs_matrix(givena,givenphi,n))
    inv = lhs_matrix(givena,n).inv()
    s = inv*rhs_v
    return vector_to_matrix(s)

def solve_n_steps(givena,phi0,n):
    print("computing q-series up to order", n, end="")
    phi = [phi0]
    for m in range(1,n+1):
        if m % 10 == 0:
            print(".", end="")
        new = solve_one_step(givena,phi,m)
        phi.append(new)
    print()
    return phi

def form_series(givena):
    size = givena[0].shape[0]
    c = zeros(size)
    for i in range(0,len(givena)):
        c += q**i*givena[i]
    return c

def check_frobenius(givena,givenphi,o):
    """
    This checks whether the Frobenius intertwining property is satisfied
    """
    quantum = form_series(givena)
    fquantum = form_series(givena).subs(q,(q**p)/(-p))
    frob = form_series(givenphi)
    intertwine = q*diff(frob,q) + quantum*frob + (-p)*frob*fquantum
    return truncate_matrix(expand(intertwine),o)

def transform_into_plot(val,v):
    """
    v is a list of valuations. We remove inf, reverse signs, rescale!
    we also remove any valuations that are larger than the promised
    accuracy
    """
    x = []
    y = []
    for i in range(0,len(v)):
        if v[i] != float("inf"):
            if v[i] <= val: 
                x.append(i)
                y.append(-v[i])
    print()
    print("Datapoints for the growth rate (p-adic valuation divided by order, for the highest 10 orders in q): ",)
    for i in range(len(x)-10,len(x)):
        print(x[i], y[i]/x[i])

    plt.margins(0)
    plt.title("p ="+str(p))
    plt.plot(x,y,".k")
    plt.plot([0,len(v)],[0,len(v)/(p-1)], color="k", linestyle = ":")
    plt.show()

def cp1(o,val):
    print("CP1")
    print()
    """
    returns a list of the valuations of the Frobenius intertwiner
    up to an including q^o, and only up to an error of val in the
    valuation
    """
    apole = Matrix([[0,0],[2,0]])
    a0 = zeros(2)
    a1 = Matrix([[0,2],[0,0]])
    givena = [apole,a0,a1]
    print("quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0],[0,Rational(1,p)]])
    alpha = eye(2)*deg
    beta = Matrix([[0,0],[1,0]])*deg
    salpha = solve_n_steps(givena,alpha,o)
    sbeta = solve_n_steps(givena,beta,o)
    malpha = min(list_of_matrix_p_adic_valuation(salpha))
    mbeta = min(list_of_matrix_p_adic_valuation(sbeta))
    m = -min(malpha,mbeta)
    gder = gamma_derivative(m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = salpha[i]+2*gder*sbeta[i]
        s.append(matrix_p_adic_valuation(new))
    return s
   
def cp2(o,val):
    print("CP2")
    print()
    apole = Matrix([[0,0,0],[3,0,0],[0,3,0]])
    a0 = zeros(3)
    a1 = zeros(3)
    a2 = Matrix([[0,0,3],[0,0,0],[0,0,0]])
    givena = [apole,a0,a1,a2]
    print("Quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0,0],[0,Rational(1,p),0],[0,0,Rational(1,p**2)]])
    alpha = eye(3)*deg
    beta = Matrix([[0,0,0],[1,0,0],[0,1,0]])*deg
    gamma = Matrix([[0,0,0],[0,0,0],[1,0,0]])*deg
    salpha = solve_n_steps(givena,alpha,o)
    sbeta = solve_n_steps(givena,beta,o)
    sgamma = solve_n_steps(givena,gamma,o)
    malpha = min(list_of_matrix_p_adic_valuation(salpha))
    mbeta = min(list_of_matrix_p_adic_valuation(sbeta))
    mgamma = min(list_of_matrix_p_adic_valuation(sgamma))
    m = -min(malpha,mbeta,mgamma)
    gder = gamma_derivative(m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = salpha[i]+3*gder*sbeta[i]+(gder**2)*Rational(9,2)*sgamma[i]
        s.append(matrix_p_adic_valuation(new))
    return s

def cp2(o,val):
    print("CP2")
    print()
    apole = Matrix([[0,0,0],[3,0,0],[0,3,0]])
    a0 = zeros(3)
    a1 = zeros(3)
    a2 = Matrix([[0,0,3],[0,0,0],[0,0,0]])
    givena = [apole,a0,a1,a2]
    print("Quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0,0],[0,Rational(1,p),0],[0,0,Rational(1,p**2)]])
    alpha = eye(3)*deg
    beta = Matrix([[0,0,0],[1,0,0],[0,1,0]])*deg
    gamma = Matrix([[0,0,0],[0,0,0],[1,0,0]])*deg
    salpha = solve_n_steps(givena,alpha,o)
    sbeta = solve_n_steps(givena,beta,o)
    sgamma = solve_n_steps(givena,gamma,o)
    malpha = min(list_of_matrix_p_adic_valuation(salpha))
    mbeta = min(list_of_matrix_p_adic_valuation(sbeta))
    mgamma = min(list_of_matrix_p_adic_valuation(sgamma))
    m = -min(malpha,mbeta,mgamma)
    gder = gamma_derivative(m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = salpha[i]+3*gder*sbeta[i]+(gder**2)*Rational(9,2)*sgamma[i]
        s.append(matrix_p_adic_valuation(new))
    return s

def cubic_surface(o,val):
    print("(Part of) the quantum connection for the cubic surface")
    print()
    apole = Matrix([[0,0,0],[1,0,0],[0,3,0]])
    a0 = Matrix([[0,0,0],[0,9,0],[0,0,0]])
    a1 = Matrix([[0,108,0],[0,0,36],[0,0,0]])
    a2 = Matrix([[0,0,252],[0,0,0],[0,0,0]])
    givena = [apole,a0,a1,a2]
    print("Quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0,0],[0,Rational(1,p),0],[0,0,Rational(1,p**2)]])
    alpha = eye(3)*deg
    beta = apole*deg
    gamma = Matrix([[0,0,0],[0,0,0],[1,0,0]])*deg
    salpha = solve_n_steps(givena,alpha,o)
    sbeta = solve_n_steps(givena,beta,o)
    sgamma = solve_n_steps(givena,gamma,o)
    malpha = min(list_of_matrix_p_adic_valuation(salpha))
    mbeta = min(list_of_matrix_p_adic_valuation(sbeta))
    mgamma = min(list_of_matrix_p_adic_valuation(sgamma))
    m = -min(malpha,mbeta,mgamma)
    gder = gamma_derivative(m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = salpha[i]+gder*sbeta[i]+(gder**2)*Rational(3,2)*sgamma[i]
        s.append(matrix_p_adic_valuation(new))
    return s

def intersection_of_quadrics(o,val):
    print("Intersection of two quadrics in P^5")
    print()
    apole = 2*Matrix([[0,0,0,0],[1,0,0,0],[0,4,0,0],[0,0,1,0]])
    a0 = zeros(4)
    a1 = 2*Matrix([[0,4,0,0],[0,0,2,0],[0,0,0,4],[0,0,0,0]])
    a2 = zeros(4)
    a3 = 2*Matrix([[0,0,0,4],[0,0,0,0],[0,0,0,0],[0,0,0,0]])
    givena = [apole,a0,a1,a2,a3]
    print("Quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0,0,0],[0,Rational(1,p),0,0],[0,0,Rational(1,p**2),0],[0,0,0,Rational(1,p**3)]])
    alpha = eye(4)*deg
    beta = Matrix([[0,0,0,0],[1,0,0,0],[0,4,0,0],[0,0,1,0]])*deg
    gamma = Matrix([[0,0,0,0],[0,0,0,0],[1,0,0,0],[0,1,0,0]])*deg
    delta = Matrix([[0,0,0,0],[0,0,0,0],[0,0,0,0],[1,0,0,0]])*deg
    salpha = solve_n_steps(givena,alpha,o)
    sbeta = solve_n_steps(givena,beta,o)
    sgamma = solve_n_steps(givena,gamma,o)
    sdelta = solve_n_steps(givena,delta,o)
    malpha = min(list_of_matrix_p_adic_valuation(salpha))
    mbeta = min(list_of_matrix_p_adic_valuation(sbeta))
    mgamma = min(list_of_matrix_p_adic_valuation(sgamma))
    mdelta = min(list_of_matrix_p_adic_valuation(sdelta))
    m = -min(malpha,mbeta,mgamma,mdelta)
    gder = gamma_derivative(m+val)
    gder3 = gamma_derivatives(3,m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = salpha[i]+2*gder*sbeta[i]+8*(gder**2)*sgamma[i]+(12*(gder**3)-Rational(20,3)*gder3)*sdelta[i]
        s.append(matrix_p_adic_valuation(new))
    return s

def twistor_space_small(o,val):
    print("Small summand of the quantum connection for the twistor space")
    print()
    apole = Matrix([[0,0,0,0],[1,0,0,0],[0,1,0,0],[0,0,1,0]])
    a0 = zeros(4)
    a1 = Matrix([[0,4,0,0],[0,0,0,0],[0,0,0,4],[0,0,0,0]])
    givena = [apole,a0,a1]
    print("Quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0,0,0],[0,Rational(1,p),0,0],[0,0,Rational(1,p**2),0],[0,0,0,Rational(1,p**3)]])
    alpha = eye(4)*deg
    beta = Matrix([[0,0,0,0],[1,0,0,0],[0,1,0,0],[0,0,1,0]])*deg
    gamma = Matrix([[0,0,0,0],[0,0,0,0],[1,0,0,0],[0,1,0,0]])*deg
    delta = Matrix([[0,0,0,0],[0,0,0,0],[0,0,0,0],[1,0,0,0]])*deg
    salpha = solve_n_steps(givena,alpha,o)
    sbeta = solve_n_steps(givena,beta,o)
    sgamma = solve_n_steps(givena,gamma,o)
    sdelta = solve_n_steps(givena,delta,o)
    malpha = min(list_of_matrix_p_adic_valuation(salpha))
    mbeta = min(list_of_matrix_p_adic_valuation(sbeta))
    mgamma = min(list_of_matrix_p_adic_valuation(sgamma))
    mdelta = min(list_of_matrix_p_adic_valuation(sdelta))
    m = -min(malpha,mbeta,mgamma,mdelta)
    gder = gamma_derivative(m+val)
    gder3 = gamma_derivatives(3,m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = salpha[i]+gder*sbeta[i]+gder**2*Rational(1,2)*sgamma[i]+gder3*Rational(1,6)*sdelta[i]
        s.append(matrix_p_adic_valuation(new))
    return s

def twistor_space_big(o,val):
    print("Big summand of the quantum connection for the twistor space")
    print()
    apole = Matrix([[0,0,0,0,0,0,0,0],[1,0,0,0,0,0,0,0],[0,1,0,0,0,0,0,0],[0,0,1,0,0,0,0,0],\
                    [0,0,0,0,0,0,0,0],[0,0,0,8,1,0,0,0],[0,0,0,0,0,1,0,0],[0,0,0,0,0,0,1,0]])
    a0 = zeros(8)
    a1 = Matrix([[0,4,0,0,0,0,0,0],[0,0,0,0,0,0,0,0],[0,0,0,4,0,0,0,0],[0,0,0,0,0,0,0,0],\
                 [0,0,0,0,0,4,0,0],[0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,4],[0,0,0,0,0,0,0,0]])
    givena = [apole,a0,a1]
    print("Quantum connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    deg = Matrix([[1,0,0,0,0,0,0,0],[0,Rational(1,p),0,0,0,0,0,0],[0,0,Rational(1,p**2),0,0,0,0,0],\
                  [0,0,0,Rational(1,p**3),0,0,0,0],[0,0,0,0,Rational(1,p**3),0,0,0],\
                  [0,0,0,0,0,Rational(1,p**4),0,0],[0,0,0,0,0,0,Rational(1,p**5),0],\
                  [0,0,0,0,0,0,0,Rational(1,p**6)]])
    euler = Matrix([[0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0],\
                [1,0,0,0,0,0,0,0],[0,1,0,0,0,0,0,0],[0,0,1,0,0,0,0,0],[0,0,0,1,0,0,0,0]])
    ch1 = apole
    ch3 = apole**3*Rational(1,6) + euler
    ch5 = apole**2*euler*Rational(7,120)
    s0 = solve_n_steps(givena,eye(8)*deg,o)
    s1 = solve_n_steps(givena,ch1*deg,o)
    s11 = solve_n_steps(givena,(ch1**2)*deg,o)
    s111 = solve_n_steps(givena,(ch1**3)*deg,o)
    s1111 = solve_n_steps(givena,(ch1**4)*deg,o)
    s11111 = solve_n_steps(givena,(ch1**5)*deg,o)
    s111111 = solve_n_steps(givena,(ch1**6)*deg,o)
    s3 = solve_n_steps(givena,ch3*deg,o)
    s31 = solve_n_steps(givena,ch3*ch1*deg,o)
    s311 = solve_n_steps(givena,ch3*(ch1**2)*deg,o)
    s3111 = solve_n_steps(givena, ch3*(ch1**3)*deg,o)
    s33 = solve_n_steps(givena,(ch3**2)*deg,o)
    s5 = solve_n_steps(givena,ch5*deg,o)
    s51 = solve_n_steps(givena,ch5*ch1*deg,o)
    stotal = s0+s1+s11+s111+s1111+s11111+s111111+s3+s31+s311+s3111+s33+s5+s51
    m = -min(list_of_matrix_p_adic_valuation(stotal))
    gder = gamma_derivative(m+val)
    gder3 = gamma_derivatives(3,m+val)
    gder5 = gamma_derivatives(5,m+val)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        new = s0[i] + gder*s1[i] + Rational(1,2)*gder**2*s11[i] + Rational(1,6)*gder**3*s111[i] \
              + Rational(1,24)*gder**4*s1111[i] + Rational(1,120)*gder**5*s11111[i] + Rational(1,720)*gder**6*s111111[i] \
              + (gder3 - gder**3)*s3[i] + (gder3 - gder**3)**2*Rational(1,2)*s33[i] \
              + (gder3 - gder**3)*gder*s31[i] + (gder3 - gder**3)*gder**2*Rational(1,2)*s311[i] \
              + (gder3 - gder**3)*gder**3*Rational(1,6)*s3111[i] \
              + (5*gder5 - 10*gder**2*gder3+9*gder**5)*s5[i] + (5*gder5 - 10*gder**2*gder3+9*gder**5)*gder*s51[i]
        s.append(matrix_p_adic_valuation(new))
    return s

def fictitious(o,val):
    print("Fictitious example of differential equation")
    print()
    apole = Matrix([[0,0,0],[1,0,0],[0,1,0]])
    a0 = Matrix([[0,0,1],[0,0,0],[0,0,0]])
    givena = [apole,a0]
    print("connection multiplied by q:")
    print()
    pprint(form_series(givena))
    print()
    alpha = eye(3)
    salpha = solve_n_steps(givena,alpha,o)
    s = []
    print("computing the p-adic valuations of",o+1,"terms in the Frobenius", end="")
    for i in range(0,o+1):
        if i % 10 == 0:
            print(".", end="")
        s.append(matrix_p_adic_valuation(salpha[i]))
    return s
  
global p
p = 5
val = 120
o = 1000
print("Prime",p)
print("Computation is exact up to p-adic valuation",val)
print("Frobenius computed up to order",o,"in the variable q")
print()
"""
we promise exactness of the computation up to p-adic valuation val. Higher val
means potentially more precision (even though only the first nonzero term counts),
but longer computation
"""
outcome = cp2(o,val)
transform_into_plot(o,outcome)