import numpy as np
import matplotlib.pyplot as plt
from numba import jit

from scvt import topology
from scvt import sphere
from scvt import delaunay
from scvt import lloyd

def test(a,b,c):
    p=circum(a,b,c)
    print(norm(a-p),norm(b-p),norm(c-p)) # should be equal

@jit
def weight(xyz):
    y = xyz[:,1];
    return 1.+15.*np.exp(-16.*y*y)

def main(nbp,inner,outer,filename=None):
    N=10*nbp*nbp+2
    points=np.random.randn(N,3)
    for i in range(outer):
        verbose = True
        print('Outer Lloyd iteration %d/%d'%(i,outer))
        for j in range(inner):
            points, vertices, edges = lloyd.lloyd(points, weight, verbose)
            verbose = False
    plt.figure(figsize=(10,10))
    sphere.plot_mesh(vertices, topology.dual_edges(edges), orient, "k")
    if N<1000:
        sphere.plot_mesh(points,edges,orient,"r")
    if filename is None:
        plt.show()
    else:
        plt.savefig(filename)



orient=np.asarray([0,0,1.])

################## Manual tests with octahedron ####################

points=np.asarray([[1.,0.,0.],[0.,1.,0.],[0.,0.,1.],
                   [-1.,0.,0.],[0.,-1.,0.],[0.,0.,-1.]])
simplices = delaunay.delaunay(points)
print("== Simplices ==", simplices)
edges, regions = topology.all_edges_regions(simplices)
print('%d cells, %d edges, %d triangles'%(points.shape[0], len(edges), len(simplices)))
print(edges)
for n,region in enumerate(regions):
    print("Voronoi cell %d :"%n, region)
for n,region in enumerate(regions):
    print("Voronoi cell %d :"%n, [edges[edge] for edge in region])
 
vertices = [sphere.circum(points[a,:],points[b,:],points[c,:]) for a,b,c in simplices]
vertices = np.asarray(vertices)
print("== Vertices ==", vertices)
bary = sphere.barycenters(regions, points, edges, vertices)

regions_size, vert = lloyd.flatten_regions(regions, edges)
degree, counts = np.unique(regions_size, return_counts=True)
print("Edge counts of Voronoi cells :", dict(zip(degree, counts)))

####################### Larger grid #######################
#main(8, 1,1)
#main(8, 100,5)
main(16, 100,5, 'nbp16.png')
main(40, 100,10, 'nbp40.png')