import numpy as np
import triangle as tr
from scipy.spatial import Delaunay as delaunay
from shapely.geometry import Polygon as ShapelyPolygon
from shapely.geometry import LineString as ShapelyLine
from shapely.geometry import Point as ShapelyPoint

class Mesh:
    def __init__(self, x, y, conmat):
        self.x = x
        self.y = y
        self.conmat = conmat

    def getNeiElements(self):
        nn = len(self.x)
        ne = len(self.conmat[:,0])
        neiElements = []
        for i in range(nn):
            ele = []
            for j in range(ne):
                if self.conmat[j,0] == i or self.conmat[j,1] == i or self.conmat[j,2] == i:
                    ele.append(j)
            neiElements.append(ele)
        self.neiElements = neiElements

    def getNeiNodes(self):
        nn = len(self.x)
        neiNodes = []
        for i in range(nn):
            nod = []
            for j in self.neiElements[i]:
                for k in self.conmat[j]:
                    if k != i: 
                        nod.append(k)
            nod = np.unique(np.array(nod))
            nod = nod.tolist()
            neiNodes.append(nod)
        self.neiNodes = neiNodes
        
def prepareGeometryForMeshing(shapes):
    # Finding all child shapes
    for i in range(len(shapes)): # for all shapes,
        if shapes[i].isClosed: # if the current shape is closed, it may contain a child.
            for j in range(len(shapes)): # for all potential child shapes,
                if i != j: # skip current shape as potential child.
                    polygon = ShapelyPolygon(shapes[i].boundaryNodes) # consider current shape as polygon,
                    if shapes[j].shape == "point":
                        point = ShapelyPoint(shapes[j].boundaryNodes) # and take one point form potential child shape.
                    else:
                        point = ShapelyPoint(shapes[j].boundaryNodes[0]) # and take one point form potential child shape.
                    if polygon.contains(point): # if one point of potential child shape is inside current shape,
                        shapes[i].child.append(j) # potential child shape is a child of current shape. Note: all the shapes must be non-contacting.

    # Finding all parent shapes
    for i in range(len(shapes)): # for all shapes,
        for j in range(len(shapes)): # for all potential parent shapes,
            if i != j: # skip current shape as potential parent.
                if shapes[j].isClosed: # if the potential shape is closed, it may be a parent.
                    polygon = ShapelyPolygon(shapes[j].boundaryNodes) # consider potential parent shape as polygon,
                    if shapes[i].shape == "point":
                        point = ShapelyPoint(shapes[i].boundaryNodes) # and take one point from current shape.
                    else:
                        point = ShapelyPoint(shapes[i].boundaryNodes[0]) # and take one point from current shape.
                    if polygon.contains(point): # if one point of current shape is inside the potential parent shape,
                        shapes[i].parent.append(j) # potential parent shape is the parent of current shape. Note: all the shapes must be non-contacting.

    for i in range(len(shapes)):
        shapes[i].parent = np.array(shapes[i].parent)
        shapes[i].child = np.array(shapes[i].child)

    # Finding immediate child shapes
    listOfShapesToRemove = []
    for i in range(len(shapes)):
        shapesToRemove = []
        for j in range(len(shapes[i].child)):
            for k in range(len(shapes[shapes[i].child[j]].child)):
                shapesToRemove.append(shapes[shapes[i].child[j]].child[k])
        listOfShapesToRemove.append(shapesToRemove)

    for i in range(len(shapes)):
        for j in range(len(listOfShapesToRemove[i])):
            shapes[i].child = np.delete(shapes[i].child, np.where(shapes[i].child == listOfShapesToRemove[i][j]))#, axis=0)

    # Finding immediate parent shapes
    listOfShapesToRemove = []
    for i in range(len(shapes)):
        shapesToRemove = []
        for j in range(len(shapes[i].parent)):
            for k in range(len(shapes[shapes[i].parent[j]].parent)):
                shapesToRemove.append(shapes[shapes[i].parent[j]].parent[k])
        listOfShapesToRemove.append(shapesToRemove)
        
    for i in range(len(shapes)):
        for j in range(len(listOfShapesToRemove[i])):
            shapes[i].parent = np.delete(shapes[i].parent, np.where(shapes[i].parent ==listOfShapesToRemove[i][j]))#, axis=0)

    # Labelling if a shape is child or parent
    for i in range(len(shapes)):
        if len(shapes[i].child) == 0:
            shapes[i].isParent = False
        else:
            shapes[i].isParent = True

    for i in range(len(shapes)):
        if len(shapes[i].parent) == 0:
            shapes[i].isChild = False
        else:
            shapes[i].isChild = True

def generateNodes(shapes, meshingMethod, maxAngle, maxArea, crse):
    # Generate interior free nodes
    if meshingMethod == "ruppert":
        for i in range(len(shapes)):
            if shapes[i].isClosed:
                meshShapeRuppert(i, shapes, crse)
            else:
                shapes[i].freeNodes = np.array([])
    elif meshingMethod == "triangle":
        meshShapesTriangle(shapes,maxAngle, maxArea)

def meshShapesTriangle(shapes, maxAngle, maxArea):
    iv, fv = triangleMesh(shapes, maxAngle, maxArea)
    initNodes = iv.shape[0]
    finalNodes = fv.shape[0]
    for shape in shapes:
        shape.freeNodes = np.empty((0,2))
    for i in range(initNodes,finalNodes):
        extraVertex = ShapelyPoint(fv[i][0],fv[i][1])
        for shape in shapes:
            if shape.isClosed:
                shapePolygon = ShapelyPolygon(shape.boundaryNodes.tolist())
                # if extraVertex is on boundary
                if not shapePolygon.contains(extraVertex) and shapePolygon.covers(extraVertex):
                    shape.boundaryNodes = np.vstack([shape.boundaryNodes, np.array([fv[i][0],fv[i][1]])])
                # if extraVertex is inside
                if shapePolygon.contains(extraVertex):
                    vertexOutside = True
                    for j in shape.child:
                        childShape = shapes[j]
                        childPolygon = ShapelyPolygon(childShape.boundaryNodes.tolist())
                        if childPolygon.covers(extraVertex):
                            vertexOutside = False
                            break
                    if vertexOutside:
                        shape.freeNodes = np.vstack([shape.freeNodes, np.array([fv[i][0],fv[i][1]])])
            elif shape.shape != 'point':
                lineString = ShapelyLine(shape.boundaryNodes.tolist())
                if lineString.contains(extraVertex):
                    shape.boundaryNodes = np.vstack([shape.boundaryNodes, np.array([fv[i][0],fv[i][1]])])

def triangleMesh(shapes, maxAngle = "", maxArea = ""):
    nov = 0
    ver = np.empty((0,2))
    seg = np.empty((0,2), dtype=int)
    for shape in shapes:
        ver = np.vstack([ver, shape.boundaryNodes])
        if shape.shape != 'point':
            seg = np.vstack([seg, np.asarray(shape.segments, dtype=int)+ nov])
        nov = ver.shape[0]
    pslg = dict(vertices=ver, segments=seg)
    stri = "pq"+maxAngle+"a"+maxArea+"A"
    tri = tr.triangulate(pslg, stri)
    return ver, tri['vertices']

def meshShapeRuppert(ind, shapes, crse): # Consider skipping sengment initialisation in case of Point
    # Parent points
    parentPoints = shapes[ind].boundaryNodes
    # Child points
    childPoints = []
    for i in range(len(shapes[ind].child)):
        childPoints.append(shapes[shapes[ind].child[i]].boundaryNodes.tolist())
    
    # Parent polygon
    parentPolygon = ShapelyPolygon(parentPoints)
    # Child polygons
    childPolygons = []
    for i in range(len(shapes[ind].child)):
        if shapes[shapes[ind].child[i]].isClosed:
            childPolygons.append(ShapelyPolygon(childPoints[i]))

    # Parent segments
    segments = []
    for i in range(len(parentPoints)-1):
        segments.append([i, i+1])
    segments.append([len(parentPoints)-1, 0])
    # Child segments
    segInd = len(parentPoints)
    for i in range(len(shapes[ind].child)):
        for j in range(len(childPoints[i])-1):
            if shapes[shapes[ind].child[i]].shape != "point":
                segments.append([segInd+j, segInd+j+1])
        if shapes[shapes[ind].child[i]].isClosed:
            segments.append([segInd+j+1, segInd])
        if shapes[shapes[ind].child[i]].shape == 'point':
            segInd = segInd + 1
        else:
            segInd = segInd + len(childPoints[i])

    # Points
    points = parentPoints

    for i in range(len(childPoints)):
        points = np.vstack((points,childPoints[i]))

    dt = delaunay(points)
    cm = dt.simplices  # Connectivity matrix
    x = points[:, 0]
    y = points[:, 1]
    ix = x.copy()
    iy = y.copy()

    elementsToRemove = []
    for i in range(len(cm)):
        cot = ShapelyPoint((x[cm[i][0]]+x[cm[i][1]]+x[cm[i][2]])/3, (y[cm[i][0]]+y[cm[i][1]]+y[cm[i][2]])/3)
        if not parentPolygon.contains(cot):
            elementsToRemove.append(i)
        for childPolygon in childPolygons:
            if childPolygon.contains(cot):
                elementsToRemove.append(i)
    cm = np.delete(cm, elementsToRemove, 0)
    st = getSharpTriangles(x, y, cm, crse)

    while len(st) != 0:
        for element in st:
            ex = x[cm[element]]
            ey = y[cm[element]]
            cc = circumCenter(ex, ey)

            isNotEcroached = segmentsEncroached(points, segments, cc)
            cirCen = ShapelyPoint(cc[0], cc[1])
            if parentPolygon.contains(cirCen) and isNotEcroached:
                pointOutsideChild = True
                for i in range(len(childPolygons)):
                    if childPolygons[i].contains(cirCen):
                        pointOutsideChild = False
                        break
                if pointOutsideChild:
                    points = np.vstack([points, cc])
                    dt = delaunay(points)
                    cm = dt.simplices # Connectivity matrix
                    x = points[:,0]
                    y = points[:,1]

                    elementsToRemove = []
                    for i in range(len(cm)):
                        cot = ShapelyPoint((x[cm[i][0]]+x[cm[i][1]]+x[cm[i][2]])/3, (y[cm[i][0]]+y[cm[i][1]]+y[cm[i][2]])/3)
                        if not parentPolygon.contains(cot):
                            elementsToRemove.append(i)
                        for childPolygon in childPolygons:
                            if childPolygon.contains(cot):
                                elementsToRemove.append(i)
                    cm = np.delete(cm, elementsToRemove, 0)
                    break
        stNew = getSharpTriangles(x, y, cm, crse)
        if stNew == st:
            break
        else:
            st = stNew

    fx = x[len(ix):len(x)-1]
    fy = y[len(iy):len(y)-1]
    freeNodes = []
    for i in range(len(fx)):
        freeNodes.append([fx[i],fy[i]])
    freeNodes = np.array(freeNodes)
    shapes[ind].freeNodes = freeNodes

def circumCenter(ex, ey):
    Ax, Ay, Bx, By, Cx, Cy = ex[0], ey[0], ex[1], ey[1], ex[2], ey[2]
    D = 2*(Ax*(By-Cy)+Bx*(Cy-Ay)+Cx*(Ay-By))
    cx = ((Ax*Ax+Ay*Ay)*(By-Cy)+(Bx*Bx+By*By)*(Cy-Ay)+(Cx*Cx+Cy*Cy)*(Ay-By))/D
    cy = ((Ax*Ax+Ay*Ay)*(Cx-Bx)+(Bx*Bx+By*By)*(Ax-Cx)+(Cx*Cx+Cy*Cy)*(Bx-Ax))/D
    return cx, cy

def circumRadiusToMinSideRatio(ex, ey):
    a, b, c = np.sqrt((ex[1]-ex[0])**2+(ey[1]-ey[0])**2), np.sqrt((ex[2]-ex[1])**2+(ey[2]-ey[1])**2), np.sqrt((ex[2]-ex[0])**2+(ey[2]-ey[0])**2)
    minSide = min([a, b, c])
    cr = (a*b*c)/(np.sqrt((a + b + c)*(b + c - a)*(c + a - b)*(a + b - c))) # Circumradius
    return cr/minSide

def getSharpTriangles(x, y, cm, B):
    st = []
    for i in range(len(cm)):
        element = cm[i]
        ex = x[element]
        ey = y[element]
        circumRatio = circumRadiusToMinSideRatio(ex, ey)
        if circumRatio > B:
            st.append(i)
    return st

def encroached(p1, p2, p):
    m = (np.array(p1)+np.array(p2))/2
    r = np.sqrt((p1[0]-m[0])**2+(p1[1]-m[1])**2)
    d = np.sqrt((p[0]-m[0])**2+(p[1]-m[1])**2)
    if d > r:
        return False
    else:
        return True

def segmentsEncroached(points, segments, cc):
    isNotEncroached = True
    for i in range(len(segments)):
        segmentCentre = [(points[segments[i][0]][0]+points[segments[i][1]][0])/2 ,(points[segments[i][0]][1]+points[segments[i][1]][1])/2]
        segmentRadius = np.sqrt((points[segments[i][0]][0]-points[segments[i][1]][0])**2+(points[segments[i][0]][1]-points[segments[i][1]][1])**2)/2
        distanceFromCentre = np.sqrt((cc[0]-segmentCentre[0])**2+(cc[1]-segmentCentre[1])**2)
        if distanceFromCentre < segmentRadius:
            isNotEncroached = False
            break
    return isNotEncroached

def createMesh(shapes):

    x = np.array([])
    y = np.array([])
    nn = -1

    for shape in shapes:
        if shape.shape == "point":
            shape.boundaryList = np.array([nn+1])
            nn = nn+1
            shape.freeList = np.array([])
            x =  np.concatenate((x, [shape.boundaryNodes[0]]))
            y =  np.concatenate((y, [shape.boundaryNodes[1]]))

        elif shape.shape == "line" or shape.shape == "polyline":
            shape.boundaryList = np.arange(nn+1, nn+len(shape.boundaryNodes[:, 0])+1, 1)
            nn = shape.boundaryList[-1]
            shape.freeList = np.array([])
            x =  np.concatenate((x, shape.boundaryNodes[:, 0]))
            y =  np.concatenate((y, shape.boundaryNodes[:, 1]))

        else:

            shape.boundaryList = np.arange(nn+1, nn+len(shape.boundaryNodes)+1, 1)
            nn = shape.boundaryList[-1]
            if len(shape.freeNodes) == 0:
                shape.freeList = np.array([])
            else:
                shape.freeList = np.arange(nn+1, nn+len(shape.freeNodes)+1, 1)
                nn = shape.freeList[-1]
            x =  np.concatenate((x, shape.boundaryNodes[:, 0], shape.freeNodes[:, 0]))
            y =  np.concatenate((y, shape.boundaryNodes[:, 1], shape.freeNodes[:, 1]))

    xy = np.transpose([x, y])
    dt = delaunay(xy)
    mesh = Mesh(x, y, dt.simplices)
    
    # Generate element list in each shape
    ne = len(dt.simplices[:,0])
    parentPolygons = []
    childPolygons = []
    for shape in shapes:
        shape.elementList = []
        if shape.isClosed:
            parentPolygons.append(ShapelyPolygon(shape.boundaryNodes))
            childPolygon = []
            for i in range(len(shape.child)):
                if shapes[shape.child[i]].shape == "point":
                    childPolygon.append([])
                else:
                    childPolygon.append(ShapelyPolygon(shapes[shape.child[i]].boundaryNodes))
            childPolygons.append(childPolygon)
        else:
            parentPolygons.append([])
            childPolygons.append([])

    cm = dt.simplices
    for i in range(ne):
        ex = x[cm[i]]
        ey = y[cm[i]]
        cc = circumCenter(ex, ey)
        point = ShapelyPoint(cc[0], cc[1])

        for j in range(len(shapes)):
            if shapes[j].isClosed:
                parentPolygon = parentPolygons[j]
                childPolygon = childPolygons[j]
                if parentPolygon.contains(point):
                    pointOutsideChild = True
                    for k in range(len(childPolygon)):
                        if shapes[shapes[j].child[k]].shape != 'point':
                            if childPolygon[k].contains(point):
                                pointOutsideChild = False
                                break

                    if pointOutsideChild:
                        shapes[j].elementList.append(i)

    # Elements
    mesh.getNeiElements()
    # Nodes
    mesh.getNeiNodes()

    return mesh

def assignSources(shapes, mesh):
    x = mesh.x
    y = mesh.y
    nn = len(x)
    # nodeSource is for both plotting and matrix assignment
    # elementSource may serve some purpose in different physics
    nodeSource = np.zeros((nn,1))
    ne = len(mesh.conmat[:,0])
    elementSource = np.zeros((ne,1))

    for shape in shapes:
        if not shape.isClosed:
            nodes = shape.boundaryList
            for node in nodes:
                nodeSource[node] = shape.source/len(nodes)
                elements = mesh.neiElements[node]
                for element in elements:
                    elementSource[element] = elementSource[element] + shape.source/len(elements)
        else:
            elements = shape.elementList
            for element in elements:
                ex = []
                ey = []
                for node in mesh.conmat[element]:
                    ex.append(x[node])
                    ey.append(y[node])
                p = [ey[1]-ey[2], ey[2]-ey[0], ey[0]-ey[1]]
                q = [ex[2]-ex[1], ex[0]-ex[2], ex[1]-ex[0]]
                ea = 0.5*abs(p[1]*q[2]-q[1]*p[2])
                elementSource[element] = shape.source*ea
            for element in elements:
                for node in mesh.conmat[element]:
                    nodeSource[node] = nodeSource[node] + elementSource[element]/3
    mesh.elementSource = elementSource            
    mesh.nodeSource = nodeSource
    
def assignMaterials(shapes, mesh):
    # nodeMaterial may be used for plotting
    # elementMaterial is for matrix computation
    x = mesh.x
    nn = len(x)
    nodeMaterial = np.ones((nn,1))
    ne = len(mesh.conmat[:,0])
    elementMaterial = np.ones((ne,1))

    for shape in shapes:
        if not shape.isClosed:
            nodes = shape.boundaryList
            for node in nodes:
                nodeMaterial[node] = shape.material
        else:
            nodes = np.concatenate((shape.boundaryList, shape.freeList)).astype(int)
            for node in nodes:
                nodeMaterial[node] = shape.material
        elements = shape.elementList
        for element in elements:
            elementMaterial[element] = shape.material
    mesh.nodeMaterial = nodeMaterial
    mesh.elementMaterial = elementMaterial

