import numpy as np
import itertools
from sage.geometry.polyhedron.ppl_lattice_polytope import LatticePolytope_PPL

x,y = var('x,y')
R = LaurentPolynomialRing(ZZ,'x,y')

def kinks(p):
    if p.dim() < 2:
        return []
    kinks = []
    bdy = p.traverse_boundary()
    for idx,i in enumerate(bdy):
        j = bdy[idx-1]
        k = bdy[(idx+1)%len(bdy)]
        m1 = p.vertices()[k] - p.vertices()[i]
        m2 = p.vertices()[j] - p.vertices()[i]
        m1 = np.array(m1)/gcd(m1)
        m2 = np.array(m2)/gcd(m2)
        kinks += [int(abs(np.linalg.det([m1,m2])))]
    return kinks

def area(p):
    if p.dim() < 2:
        return 0
    return len(p.interior_points()) + 1/2 * len(p.boundary_points()) - 1

def lengths(p):
    if p.dim() < 2:
        return [p.npoints()-1]
    return [e.npoints()-1 for e in p.edges()]

def size(p):
    return max([max(abs(v1[0]-v2[0]),abs(v1[1]-v2[1])) for v1,v2 in itertools.product(p.vertices(),repeat=2)])

def tex_poly(f,p,scale,N):
    x_min = min([v[0] for v in p.vertices()])
    x_max = max([v[0] for v in p.vertices()])
    y_min = min([v[1] for v in p.vertices()])
    y_max = max([v[1] for v in p.vertices()])
    width = max([1,x_max-x_min,y_max-y_min])
    output = f"\\begin{{tikzpicture}}[scale=2/{N*scale}] "
    if p.dim() == 0:
        output += "\\fill (0,0) circle (1pt); "
    elif p.dim() == 1:
        output += f"\\draw {tuple(p.vertices()[0])} -- {tuple(p.vertices()[1])}; "
    else:
        output += "\\draw " + ' -- '.join([str(tuple(p.vertices()[i])) for i in p.traverse_boundary()]) + " -- cycle; "
    for v in p.points():
        coeffs = restrict(f,[tuple(v)]).coefficients()
        coeff = coeffs[0] if len(coeffs) > 0 else 0
        output += f"\\fill {tuple(v)} node{{\\tiny${coeff}$}} circle (.5pt);"
    output += "\\end{tikzpicture}"
    return output

def tex(vertices,edges,N=1):
    vertices = [(f,newton_polytope(f)) for f in vertices]
    max_vert = max([p.nvertices() for f,p in vertices])
    x_vars = []
    y_vars = []
    count = [0 for n in range(max_vert+1)]
    for f,p in vertices:
        n = p.nvertices()
        x_vars.append(n)
        y_vars.append(count[n])
        count[n] += 1
    x_vars = [round(x/max(x_vars),3) for x in x_vars]
    y_vars = [round(y/max(y_vars),3) for y in y_vars]
    scale = max([1,2*(max(x_vars)-min(x_vars)),max(y_vars)-min(y_vars)])
    output = "\\begin{figure}[h!]\n"
    output += "\\centering\n"
    output += f"\\begin{{tikzpicture}}[xscale=25/{scale},yscale=20]\n"
    for i in range(len(vertices)):
        output += f"\\node ({i}) at ({x_vars[i]},{y_vars[i]}) {{}}; "
    output += "\n"
    for e in edges:
        output += f"\\draw[gray] ({e[0]}) -- ({e[1]}); "
    output += '\n'
    for i in range(len(vertices)):
        f,p = vertices[i]
        if i == 0:
            output += f"\\draw[red] ({i}) node {{"
        else:
            output += f"\\draw ({i}) node {{"
        output += tex_poly(f,p,scale,N)
        output += "};\n"
    output += "\\end{tikzpicture}\n"
    output += "\\caption{A mutation graph.}\n"
    output += "\\end{figure}"
    print(output)

def newton_polytope(f):
    return LatticePolytope([tuple(term[1].exponents()[0]) for term in list(R(f))])

def restrict(f,p):
    p = [tuple(v) for v in LatticePolytope(p).points()]
    return R(sum([term[0]*term[1] for term in list(R(f)) if tuple(term[1].exponents()[0]) in p]))

def mut(f,g,w,b):
    p = newton_polytope(f)
    h_min = min([np.dot(w,v)+b for v in p.vertices()])
    h_max = max([np.dot(w,v)+b for v in p.vertices()])
    res = 0
    for h in range(h_min,h_max+1):
        p_h = [tuple(v) for v in p.points() if np.dot(w,v)+b == h]
        f_h = restrict(f,p_h) if len(p_h) > 0 else 0
        if h < 0:
            if f_h/g^(-h) not in R:
                return 'not mutable'
            res += R(f_h/g^(-h))
        else:
            res += R(f_h*g^h)
    return res

def index(f,lst):
    res = (-1,f)
    p = newton_polytope(f)
    p_ppl = LatticePolytope_PPL([tuple(v) for v in p.vertices()])
    for i,g in enumerate(lst):
        q = newton_polytope(g)
        q_ppl = LatticePolytope_PPL([tuple(v) for v in q.vertices()])
        if p_ppl.is_isomorphic(q_ppl):
            iso = p_ppl.find_isomorphism(q_ppl) 
            iso_f = 0
            for v in p.points():
                iso_v = tuple(iso(LatticePolytope_PPL(tuple(v))).vertices()[0])
                coeffs = restrict(f,[tuple(v)]).coefficients()
                coeff = coeffs[0] if len(coeffs) > 0 else 0
                iso_f += coeff*x^iso_v[0]*y^iso_v[1]
            iso_f = R(iso_f)
            if R(g) == iso_f:
                return (i,iso_f)
            else:
                res = (-1,iso_f)
    return res

def add_and_next(f,f_old_lst,N,vertices,edges):
    if f == 'not mutable':
        return
    p = newton_polytope(f)
    if size(p) > N:
        return
    index_old_lst = [index(f_old,vertices)[0] for f_old in f_old_lst]
    index_new,f_iso = index(f,vertices)
    # if f is old (up to isomorphism)
    if index_new >= 0:
        for index_old in index_old_lst:
            if index_old != index_new:
                edge = (index_old,index_new) if index_old < index_new else (index_new,index_old)
                if edge not in edges:
                    edges.append(edge)
        return
    # if f is new
    vertices.append(f_iso)
    index_new = len(vertices)-1
    for index_old in index_old_lst:
        edges.append((index_old,index_new))
    print(index_old_lst,"->",index_new,"\r\t\t lengths ",lengths(p),"\r\t\t\t\t\t\t kinks ",kinks(p),"\r\t\t\t\t\t\t\t\t\t\t area ",area(p),"\t polytope ",[tuple(v) for v in p.vertices()],"\r\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t polynomial ",f_iso)
    # products and next mutation
    for f_factor in vertices:
        if f_factor != 1:
            add_and_next(f_iso*f_factor,[f_iso,f_factor],N,vertices,edges)
    mut_step(f_iso,N,vertices,edges)

def mut_step(f,N,vertices,edges):
    p = newton_polytope(f)
    # multiples
    if f == 1:
        for n in range(1,N+1):
            f_new = (1+x)^n
            add_and_next(f_new,[f],N,vertices,edges)
        return
    # mutate smaller
    for w in p.facet_normals():
        for u in [(-w[1]/gcd(w[0],w[1]),w[0]/gcd(w[0],w[1])),(w[1]/gcd(w[0],w[1]),-w[0]/gcd(w[0],w[1]))]:
            g = R(1+x^u[0]*y^u[1])
            h_min = min([np.dot(w,v) for v in p.vertices()])
            h_max = max([np.dot(w,v) for v in p.vertices()])
            slices = [[v for v in p.points() if np.dot(w,v) == h] for h in range(h_min,h_max+1)]
            # distance between graph of phi_0 and slice length
            dist = min([len(slice)-1 + h for h,slice in enumerate(slices)])
            for b in range(-dist-h_min,-dist+len(slices[0])):
                f_new = mut(f,g,w,b-h_min)
                if f_new != "not mutable":
                    p = newton_polytope(f_new)
                    if size(p) <= N:
                        print([tuple(v) for v in p.vertices()],f_new,",",f,",",g,w,b,h_min,"smaller")
                add_and_next(f_new,[f],N,vertices,edges)
    # mutate bigger
    for w in itertools.product(range(-N,N+1),repeat=2):
        if gcd(w) != 1:
            continue
        w0min = sorted(p.vertices(),key=lambda v: np.dot(w,v))[0]
        for u in [(-w[1]/gcd(w[0],w[1]),w[0]/gcd(w[0],w[1])),(w[1]/gcd(w[0],w[1]),-w[0]/gcd(w[0],w[1]))]:
            g = R(1+x^u[0]*y^u[1])
            h_min = min([np.dot(w,v) for v in p.vertices()])
            for b in range(N):
                f_new = mut(f,g,w,b-h_min)
                if f_new != "not mutable":
                    p = newton_polytope(f_new)
                    if size(p) <= N:
                        print([tuple(v) for v in p.vertices()],f_new,",",f,",",g,w,b,h_min,"bigger")
                add_and_next(f_new,[f],N,vertices,edges)

def muts(f,N):
    vertices = [f]
    edges = []
    mut_step(f,N,vertices,edges)
    edges = sorted(set(edges))
    return vertices,edges

def delete_reducible(vs,es):
    vertices = []
    removed = []
    edges = []
    for i,f in enumerate(vs):
        if sum([mult for factor,mult in R(f).factor()]) <= 1:
            vertices.append(f)
        else:
            removed.append(i)
    for e in es:
        if e[0] in removed or e[1] in removed:
            continue
        edge = (e[0]-len([i for i in removed if i < e[0]]),e[1]-len([i for i in removed if i < e[1]]))
        edges.append(edge)
    return vertices,edges

#f = (1+x)^2+x*y^2
#f = (1+x)^4+y*(1+4*x+3*x^2)+y^2*(1+2*x)+x^3
#f = (1+x)^3+2*y*(1+x)+y^2
f = 1

N = 2

vertices,edges = muts(f,N)
#vertices2,edges2 = delete_reducible(vertices,edges)
tex(vertices,edges,1/2)