import bpy
import random
import bmesh
import math
from mathutils import *

def centroid(coords):

    sum = []
    sum.extend(coords[0])
    n = len(coords)
    for i in range(1, n):
        for j in range(len(sum)):
            sum[j] += coords[i][j]
    return [ x/ n for x in sum]
    
def bridgeRings(top, bottom):
    """
    :param top: a list of vertex indices for the top ring
    :param bottom: a list of verex indices for the bottom ring
    :return: a list of faces bridging the two rings
    """

    rval  = []
    n = len(top)
    for i in range(n):
        v1 = top[i]
        v2 = top[ (i+1)%n]
        v3 = bottom[i]
        v4 = bottom[ (i+1)%n]
        rval.append( [ v1, v3, v4, v2] )
    return rval


def make_prism(name, z1, z2, coords, zBevel, insetFactor):
    center = centroid(coords)
    ring1 = [[c[0] - center[0], c[1] - center[1], z1] for c in coords]
    ring2 = [[c[0] - center[0], c[1] - center[1], z2] for c in coords]
    ring2Low = [[c[0] - center[0], c[1] - center[1], z2 - zBevel] for c in coords]
    ring2inset = [[insetFactor * (c[0] - center[0]), insetFactor * (c[1] - center[1]), z2] for c in coords]
    verts = []
    verts.extend(ring1)
    verts.extend(ring2Low)
    verts.extend(ring2)
    verts.extend(ring2inset)
    faces = []
    ringList = [
        [i + base for i in range(len(coords))] for base in range(0, 4 * len(coords), len(coords))
    ]
    # bottom = [ i for i in range(len(coords)) ]
    faces.append(list(reversed(ringList[-1])))
    #top = [ len(coords)*3 + i for i in range(len(coords)) ]
    faces.append(ringList[0])
    for i in range(1, len(ringList)):
        faces.extend(bridgeRings(ringList[i - 1], ringList[i]))
    mesh = bpy.data.meshes.new(name)
    #print(verts)
    #print(faces)
    mesh.from_pydata(verts, [], faces)
    #
    for face in mesh.polygons:
        face.use_smooth = True

    #
    obj = bpy.data.objects.new(name, mesh)
    obj.location = [center[0], center[1], 0]
    return obj


def interp_vvw(a, b, cornerBevel):
    return Vector(a) * (1 - cornerBevel) + Vector(b) * cornerBevel


def make_prism_beveled(name, z1, z2, coords, cornerBevel=0.1, zBevel=0.3, insetFactor=0.7):
    c2 = []
    for i in range(len(coords)):
        a = coords[i]
        b = coords[ (i+1)%len(coords)]

        c2.append(a)
        c2.append( interp_vvw(a, b, cornerBevel) )
        c2.append( interp_vvw(a,b, (1-cornerBevel)))

    obj = make_prism(name, z1, z2, c2, zBevel, insetFactor)

    obj.modifiers.new("subdiv", 'SUBSURF')

    return obj


def make_prism_sharp(name, z1,z2, coords, zBevel = 0.3, insetFactor = 0.7):
    obj = make_prism(name, z1, z2, coords, zBevel, insetFactor)

    #

    obj.modifiers.new("subdiv", 'SUBSURF')

    bm = bmesh.new()
    bm.from_mesh(obj.data)

    cr = bm.edges.layers.crease.active
    if cr is None:
        cr = bm.edges.layers.crease.new()

    creaseRings = set([1,2])
    for e in bm.edges:
        ringa = [ math.floor(v.index/len(coords)) for v in e.verts ]
        if ( ringa[0] != ringa[1] and (ringa[0] in creaseRings or ringa[1] in creaseRings) ):
            e[cr] = 1

    bm.to_mesh(obj.data)

    #

    obj.modifiers.new("edgesplit", 'EDGE_SPLIT')

    return obj

def prismMaterials():
    rval = []
    for mat in bpy.data.materials:
        if (mat.name[:5] == "stick"):
            rval.append(mat)
    return rval

def robinson(v1, v2, v3):
    obj = make_prism_sharp("robinson", -4, random.uniform(0,1), [v1,v2,v3], 0.1)
    scn.objects.link(obj)
    obj.layers = [i in set([ 2, 5]) for i in range(20)]

def robinsonq(v1, v2, v3, v4, acute):
    obj = make_prism_beveled("robinson", -4, random.uniform(0.5,1), [v1,v2,v3, v4], 0.04)
    scn.objects.link(obj)
    obj.layers = [i in set([ 2, 5]) for i in range(20)]
    if (acute):
        mat = bpy.data.materials.get("stick 1")
    else:
        mat = bpy.data.materials.get("stick 2")
    obj.data.materials.append(mat)

def robinsons():

    robinson([1.6653345369377348E-16,0.0], [0.8620385149429254,-0.5068427751827931], [0.39948909991113335,-0.9167379445905968])
    robinson([0.39948909991113335,-0.9167379445905968], [-0.3489315750866878,-1.5799622605168193], [1.6653345369377348E-16,0.0])
    robinson([1.6653345369377348E-16,0.0], [0.8620385149429254,-0.5068427751827931], [0.21565157315160932,0.97647037794151])
    robinson([-0.7796669437263946,0.8798212705489941], [-0.9953185168780039,-0.09664910739251598], [0.21565157315160932,0.97647037794151])
    robinson([1.6653345369377348E-16,0.0], [-0.9953185168780039,-0.09664910739251598], [-0.7484206749978211,-0.6632243159262224])
    robinson([-0.7484206749978211,-0.6632243159262224], [-0.3489315750866878,-1.5799622605168193], [1.6653345369377348E-16,0.0])
    robinson([1.6653345369377348E-16,0.0], [-0.9953185168780039,-0.09664910739251598], [0.21565157315160932,0.97647037794151])

def squares():
    global colors, u, v, obj, ci, i
    colors = prismMaterials()
    for u in range(w):
        for v in range(w):
            obj = make_prism_sharp("stick", -4, random.uniform(0, 2), [[u, v], [u + 1, v], [u + 1, v + 1], [u, v + 1]])
            scn.objects.link(obj)
            ci = random.randint(0, len(colors) - 1)
            print(ci)
            obj.data.materials.append(colors[ci])
            obj.layers = [i in set([0, 3, 5]) for i in range(20)]


#
#

scn = bpy.context.scene
w=20

            # obj = makePrism("bacon", 0, 2, [ [0,0], [1,0], [0.5, 0.7]])
            #scn.objects.link(obj)


#squares()
#robinsons()