import itertools

print('Run run_checks()')

# We fix 12 points and 12 roots for the local monodromy
pt = [k for k in range(12)]
num12 = [20, 0, 5, 9, 12, 14, 4, 34, 15, 32, 25, 24]

R.<t> = QQ[]

# First, we define 27 lines and 36 roots, as vectors in I^{1,6} 

m1_6 = diagonal_matrix([1] + [-1 for _ in range(6)])
I1_6 = VectorSpace(QQ, 7, inner_product_matrix=m1_6)

line = (
    [
        I1_6([1 if k==i else 0 for k in range(7)])
        for i in range(1, 7)
    ] +
    [
        I1_6([0 if k==i else 2 if k==0 else -1 for k in range(7)])
        for i in range(1, 7)
    ] +
    [
        I1_6([-1 if k==i or k==j else 1 if k==0 else 0 for k in range(7)])
        for i in range(1, 7) for j in range(i+1, 7)
    ]
)

root = (
    [
        I1_6([1 if k==i else -1 if k==j else 0 for k in range(7)])
        for i in range(1, 7) for j in range(i+1, 7)
    ] + 
    [
        I1_6([
            -1 if (n==i or n==j or n==k) else 1 if n==0 else 0 for n in range(7)
        ])
        for i in range(1, 7) for j in range(i+1, 7) for k in range(j+1, 7)
    ] +
    [ I1_6([2 if k==0 else -1 for k in range(7)]) ]
)
 
# Intersection numbers
def check_intersection_numbers():
    root_line = matrix([ [r.inner_product(l) for l in line] for r in root ])

    print('Squares of roots:')
    for r in root:
        print(r.inner_product(r),)
    print('\n\nSquares of lines:')
    for l in line:
        print(l.inner_product(l),)

    print('\n\nFor each root, num of lines intersecting by 1, -1, 0:')
    for k in range(len(root)):
        row = list(root_line.row(k))
        print([row.count(1), row.count(-1), row.count(0)], )
    print('\n\nFor each line, num of roots intersecting by 1, -1, 0:')
    for k in range(len(line)):
        col = list(root_line.column(k))
        print([col.count(1), col.count(-1), col.count(0)],)

root12 = [root[k] for k in num12]        

class Section():
    r"""A section of \oplus_{i=1}^{27} H^0(P^1, O(d_i))"""
    def __init__(self,
                 degs,          # A list of 27 numbers
                 polys=None,    # A list of 27 polynomials
                 vecs=None,     # A list of 27 vectors, or lists
                 longvec=None   # One long vector
                 ):
        self.degs = degs
        self.lens = [deg + 1 for deg in degs] # lengths of vecs

        # One of these is required
        if polys != None:
            self.polys = polys
            self.vecs = [poly_to_vec(polys[k], degs[k]) for k in range(27)]
            self.longvec = flat_list(self.vecs)
        elif vecs != None:
            self.vecs = vecs
            self.longvec = flat_list(self.vecs)
            self.polys = [vec_to_poly(vec) for vec in vecs]
        elif longvec != None:
            self.longvec = longvec
            self.vecs = [
                longvec[
                    sum(self.lens[0 : k]) : sum(self.lens[0 : (k+1)])
                ] for k in range(27)
            ]
            self.polys = [vec_to_poly(vec) for vec in self.vecs]

    def __str__(self):
        return 'degs = ' + str(self.degs) + ' polys = ' + str(self.polys)

    def __repr__(self):
        return str(self)

    def __hash__(self):
        return hash(str(self))

    def __mul__(self, other):
        degs = [self.degs[k] + other.degs[k] for k in range(27)]
        polys = [self.polys[k] * other.polys[k] for k in range(27)]
        return Section(degs=degs, polys=polys)        
            
def poly_to_vec(poly, deg):
    try:
        return [poly.constant_coefficient()] + [poly.monomial_coefficient(t^n) for n in range(deg + 1)]
    except:  # if poly is a constant
        return [poly] + [0  for n in range(deg + 1)]

def vec_to_poly(vec):
    return sum([vec[k] * t^k for k in range(len(vec))])

def flat_list(L):
    return list(itertools.chain.from_iterable([list(x) for x in L]))

# The glue for gluing 27 P^1's: a list of 12 lists, each of them is a list
# of ordered pairs of sheets to glue

glue = []

for r in root12:
    ith_glue = []
    for il, l in enumerate(line):
        if r.inner_product(l) == 1: 
            l2 = l + r      # reflected line
            il2 = line.index(l2)
            ith_glue.append([il, il2])
    glue.append(ith_glue)

def polys_upto_degs(degs):
    lens = [d + 1 for d in degs]
    r"""Return \oplus_1^27 [1, t, ... t^d]"""
    return [
        [t^k if s == l else 0 * t^0 for s in range(27)]
        for l in range(27) for k in range(lens[l])
    ]

def product(L):
    """Product of all elements in a list L"""
    return reduce(lambda x, y: x*y, L)

def glue_sections(degs, Lsecs, glue, mK):
    """Input is a basis in a vector space of Sections. Output is sections that glue with at the sheets specified by the glue. Gluing is with a twist corresponding to the integers mK"""
    lens = [d + 1 for d in degs]
    L = []
    for i in range(12):
        for pair in glue[i]:
            [l1, l2] = pair
            row = [
                sec[l1](pt[i]) /
                (product([pt[i] - pt[k] for k in nodes_C[l1] if k != i]))^mK
                -
                sec[l2](pt[i]) /
                (-product([pt[i] - pt[k] for k in nodes_C[l2] if k != i]))^mK
                for sec in Lsecs
            ]
            # We divided by the product because e.g. a section of H^0(K)
            # is identified with a collection of polynomials P(t), and
            # P(t) corresponds to P(t) / \prod_{nodes on sheet l} (t-pt[k]
            # which has a residue 1 / \prod_{nodes on sheet l!=i} (pt[i]-pt[k]
            # at the point pt[i]
            L.append(row)
    m = matrix.column(L)
    new_basis = m.kernel().basis()
    return [Section(degs=degs, polys = [sum([b[k] * Lsecs[k][l] for k in range(len(b))]) for l in range(27)]) for b in new_basis]

def h0dL(d):
    return len(H0dL(d))

def H0dL(d):
    degs = [d for _ in range(27)]
    Lsecs = polys_upto_degs(degs)
    return glue_sections(degs, Lsecs, glue, 0)

# Degrees of KC restricted to the 27 irreducible components of C
def count_pm1(x):
    L = list(x)
    return L.count(1) + L.count(-1)

def places_pm1(x):
    L = list(x)
    return [k for k in range(len(L)) if L[k]==1 or L[k]==-1]

# A 27-list of lists of numbers where k-th curve has nodes
nodes_C = [ 
    places_pm1([l.inner_product(r) for r in root12]) 
    for l in line
]

degs_KC = [
    count_pm1([l.inner_product(r) for r in root12]) - 2
    for l in line
]

def H0mkdl(m, d):
    """Returns a basis of H^0(mK + dL)"""
    degs = [m * x + d for x in degs_KC]
    Lsecs = polys_upto_degs(degs)
    return glue_sections(degs, Lsecs, glue, m)

def h0mkdl(m, d):
    return len(H0mkdl(m, d))

def multiply_spaces(V1, V2):
    if len(V1) == 0 or len(V2) == 0:
        return []
    degs = [V1[0].degs[k] + V2[0].degs[k] for k in range(27)]
    longvecs = [(s1 * s2).longvec for s1 in V1 for s2 in V2]
    m = matrix(longvecs)
    L = m.image().basis()
    return [Section(degs, longvec=row) for row in L]

# The space H^0(K)^{-5} is the subspace of H^0(K) where all 6 residues concide 

def H0k_5():
    """Return a basis of H^0(K)^{-5}"""
    Lsecs = H0mkdl(1, 0)
    lens = [d + 1 for d in degs_KC]
    L = []
    for i in range(12):
        ls = [pair[0] for pair in glue[i]]
        l1 = ls[0]
        for l2 in ls[1:]:
            row = [
                sec.polys[l1](pt[i]) /
                (product([pt[i] - pt[k] for k in nodes_C[l1] if k != i]))
                -
                sec.polys[l2](pt[i]) /
                (product([pt[i] - pt[k] for k in nodes_C[l2] if k != i]))
                for sec in Lsecs
            ]
            # We divided by the product because e.g. a section of H^0(K)
            # is identified with a collection of polynomials P(t), and
            # P(t) corresponds to P(t) / \prod_{nodes on sheet l} (t-pt[k]
            # which has a residue 1 / \prod_{nodes on sheet l!=i} (pt[i]-pt[k]
            # at the point pt[i]
            L.append(row)
    m = matrix.column(L)
    new_basis = m.kernel().basis()
    return [Section(degs=degs_KC, polys = [sum([b[k] * Lsecs[k].polys[l] for k in range(len(b))]) for l in range(27)]) for b in new_basis]


def run_checks():
    print('Chosen points:', pt)
    print('Chosen roots:')
    for k, r in enumerate(root12):
        print(num12[k], '\t', r)

    V = H0k_5()
    if len(V) != 6:
        print('Error!')
        return
    else:
        print('\nLet V = H^0(K_C)^{-5}. We have dim V =', len(V))
        I = [R.ideal([sec.polys[l] for sec in V]) for l in range(27)]
        Itrivial = R.ideal([1])
        if not all([ J == Itrivial for J in I]):
            print('There are base points of V on A^1')
            return
        if not [
                max([sec.polys[l].degree() for sec in V])
                for l in range(27)
        ] == degs_KC:
            print('There is a base point of V at infinity')
            return
        print('V is base point free')
        
        W = multiply_spaces(V, V)
        if len(W) != 21:
            print('Error!')
            return
        else:
            print('\nLet W = im (V * V) in H^0(2K)')
            print('We have dim W = 21, so the Prym-Torelli theorem does not lie on a quadric')
    
        U1 = H0mkdl(0, 1)
        U2 = H0mkdl(1, -1)
        U3 = multiply_spaces(U1, U2)
        if not (len(U1) == 2 and len(U2) == 20 and len(U3) == 40):
            print('Error!')
            return
        else:
            print('\nLet U1 = H^0(L), U2 = H^0(K-L), U3 = im (U1 * U2) in H^0(K)')
            print('We have dim U1 =', len(U1), 'dim U2 =', len(U2), 'dim U3 =', len(U3))
            print('Thus, the Petri map is generically injective')