#!/usr/bin/python3

import numpy as np
from scipy.spatial import Delaunay
import matplotlib.pyplot as plt
import shapely.geometry as geo

#import sys
#np.set_printoptions(threshold=sys.maxsize)

def get_mesh(fd,fh,h0,bbox,pfix):
    '''
    A simple mesh generator based on:
    Persson, G. Strang, A Simple Mesh Generator in MATLAB.
    SIAM Review, Volume 46 (2), pp. 329-345, June 2004
    '''
    geps = 0.001*h0
    ttol = 0.1
    Fscale = 1.2
    deltat = 0.2
    deps = np.sqrt(np.finfo(float).eps)*h0
    dptol = 0.001
    maxIter = 500

    X, Y = np.meshgrid(np.arange(bbox[0,0],bbox[1,0]+h0,h0),\
            np.arange(bbox[0,1],bbox[1,1]+h0*0.5*np.sqrt(3),h0*0.5*np.sqrt(3)))
    X[1:-1:2,:] += (0.5*h0)
    p = np.column_stack((X.flatten('F'),Y.flatten('F')))

    p = p[np.nonzero(np.apply_along_axis(fd,1,p) < geps)].squeeze()
    r0 = 1/np.apply_along_axis(fh,1,p)**2
    R0 = np.max(r0)
    p = p[np.random.rand(p.shape[0]) < r0/R0]
    if len(pfix):
        p = np.vstack((pfix, p))
    N = p.shape[0]

    pold = np.Inf
    counter = 0
    while True:
        if np.max(np.sqrt(np.sum((p-pold)**2,1))/h0) > ttol:
            pold = np.copy(p)
            T = Delaunay(p)
            tri = T.simplices
            pmid = (p[tri[:,0],:]+p[tri[:,1],:]+p[tri[:,2],:])/3
            inner = np.apply_along_axis(fd,1,pmid) < -geps
            tri = tri[inner,:]
            bars = np.vstack((tri[:,[0,1]],tri[:,[1,2]],tri[:,[0,2]]))
            bars.sort(1)
            bars = remove_duplicate_row(bars)


        bvec = p[bars[:,0],:] - p[bars[:,1],:]
        L = np.sqrt(np.sum(bvec**2,1))
        hbars = np.apply_along_axis(fh, 1,(p[bars[:,0],:]+p[bars[:,1],:])/2)
        L0 = hbars*Fscale*np.sqrt(np.sum(L**2)/np.sum(hbars**2))
        F = [np.max(L0 - [L], 0)]
        Fvec = (F/L).T * bvec 

        #sparse and back to full does not work the same in python as in matlab
        #forced to do the loop, figure out better way
        row = bars[:,[0,0,1,1]].flatten('F')
        column = (np.ones(np.shape(F),dtype=np.int32).T @ [[0,1,0,1]]).flatten('F')
        Ffrac = np.hstack((Fvec,-Fvec)).flatten('F')
        Ftot = np.zeros([N,2])
        for i in range(len(row)):
            Ftot[row[i],column[i]] += Ffrac[i]
        if len(pfix):
            Ftot[0:pfix.shape[0]-1,:] = 0
        p += (deltat*Ftot)

        #bring back to boundary
        d = np.apply_along_axis(fd,1,p)
        ix = d>0 #index of points outside boundary
        dx = np.column_stack((p[ix,0]+deps, p[ix,1]))
        dy = np.column_stack((p[ix,0], p[ix,1]+deps))

        gradx = (np.apply_along_axis(fd,1,dx)-d[ix])/deps
        grady = (np.apply_along_axis(fd,1,dy)-d[ix])/deps
        p[ix,:] -= np.column_stack((d[ix]*gradx,d[ix]*grady))

        #print(np.max(np.sqrt(np.sum(deltat*Ftot[d<-geps,:]**2,1))/h0))

        counter += 1
        #print(counter)

        if Ftot[d<-geps,:].size == 0 or counter > maxIter:
            print(counter)
            return p, tri
        if np.max(np.sqrt(np.sum(deltat*Ftot[d<-geps,:]**2,1))/h0) < dptol:
            print(counter)
            return p, tri

def remove_duplicate_row(M):
    #remove dupplicate rows of the array
    A = np.ascontiguousarray(M).view(np.dtype((np.void,M.dtype.itemsize*M.shape[1])))
    _, idx = np.unique(A, return_index=True)
    return M[idx]

def print_mesh(p, tri, title='', output=''):
    plt.figure()
    plt.triplot(p[:,0], p[:,1], tri, color='black')
    plt.axis('equal')
    plt.axis(False)
    if title:
        plt.title(title)
    if output:
        plt.savefig(output)
        plt.close()
        return
    plt.show()
    return

def ddiff(d1, d2):
    return max(d1, -d2)

def dunion(d1, d2):
    return min(d1, d2)

def dintersect(d1, d2):
    return max(d1, d2)

def dcircle(p, xc, yc, r):
    return np.sqrt((p[0]-xc)**2+(p[1]-yc)**2)-r

def drectangle(p, x1, x2, y1, y2):
    return -min(min(min(-y1+p[1],y2-p[1]),-x1+p[0]), x2-p[0])

def protate(p, phi):
    A = np.array([[np.cos(phi),-np.sin(phi)],[np.sin(phi),np.cos(phi)]])
    return (A@p.T).T

def _dpoly(p, poly, ring):
    ds = ring.distance(p)
    inside = ring.contains(p) or poly.contains(p)
    return (-1)**inside*ds

def dpoly(p, poly):
    return _dpoly(geo.asPoint(p), geo.asPolygon(poly),geo.asLinearRing(poly))

if __name__ == "__main__":
    #print("test 1")
    #fd = lambda x: ddiff(dcircle(x, 0, 0, 1), dcircle(x, 0.5, -0.2, 0.2))
    #fh = lambda x: 0.8
    #bbox = np.array([[-1, -1],[1, 1]])
    #p, tri = get_mesh(fd, fh, 0.1, bbox, [])
    #print_mesh(p, tri)

    print("test 2")
    import time
    poly = np.array([[0,0],[0,1],[1,1],[1,2],[2,2],[2,1],[3,1],[3,0]])
    polyG = geo.asPolygon(poly)
    ringG = geo.asLinearRing(poly)
    fd = lambda x: dpoly(x,poly)
    fh = lambda x: 0.15-0.22*fd(x)
    bbox = np.array([[0,0],[3,3]])
    t = time.time()
    p, tri = get_mesh(fd, fh, 0.1, bbox, poly)
    elap = time.time()-t
    print("prvi = {}s".format(elap))
    fd = lambda x: _dpoly(geo.asPoint(x), polyG, ringG)
    fh = lambda x: 0.15-0.22*fd(x)
    t = time.time()
    p, tri = get_mesh(fd, fh, 0.1, bbox, poly)
    elap = time.time()-t
    print("drugi = {}s".format(elap))
    print_mesh(p, tri)
    

    #print("test 3")
    #fd = lambda x: drectangle(protate(x,np.pi/4),0,1,0,1)
    #bbox = np.array([[-2,-2],[2,2]])
    #p, tri = get_mesh(fd, fh, 0.08, bbox, [])
    #print_mesh(p, tri)