import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse.linalg import eigs
import time
from scipy.sparse import dok_matrix
from scipy.spatial import Delaunay
from matplotlib.tri import Triangulation

# Mesh construction
r = np.linspace(0, 1, 31)  # radial coordinates
dxt = r[1]-r[0]
xc = [0.0]
yc = [0.0]
# Put in nodes circle by circle
for i in range(1,len(r)):
    Np = int((2*np.pi*r[i]) // dxt)  # number of points on this circle
    phi = np.linspace(0, 2*np.pi, Np+1)
    xt = r[i]*np.cos(phi[0:-1])
    yt = r[i]*np.sin(phi[0:-1])
    xc += list(xt)
    yc += list(yt)
x = np.array(xc)
y = np.array(yc)

# Plot just the points
plt.figure()
plt.plot(x,y,'o')
#plt.show()
#exit()

# Sort nodes; may or may not help performance
i = np.argsort(x)
x = x[i]
y = y[i]
i = np.argsort(y)
x = x[i]
y = y[i]

# Delaunay triangulation
c = np.stack([x.flat, y.flat], axis=1)  # N by 2
tri = Delaunay(c)
Ne = tri.simplices.shape[0] # number of elements
centers = tri.points[tri.simplices].mean(axis=1) # N by 2
N = tri.points.shape[0] # number of nodes

# Plot the mesh
plt.figure()
plt.triplot(x,y,tri.simplices)
plt.axis('equal')

# permittivity on a element-by-element basis.
eps = np.ones(Ne)
for ie in range(Ne):
    if centers[ie,0]**2 + centers[ie,1]**2 <= 0.3**2:
        eps[ie] = 12

k0 = 2*np.pi

laplace = dok_matrix((N,N))
wavenumop = dok_matrix((N,N))
mass = dok_matrix((N, N))

# Prototype basis function gradients
gradp2 = np.array([-1,-1])
gradp0 = np.array([1,0])
gradp1 = np.array([0,1])

t1 = time.time()
for ie in range(Ne):
    # Get the transformation.
    A = tri.transform[ie,0:2,0:2]
    invA = np.linalg.inv(A)
    detinvA = invA[0,0]*invA[1,1] - invA[1,0]*invA[0,1]
    # Indices of nodes that touch this element
    ip = tri.simplices[ie,:]
    # Note that the barycentric coordinate system is set up such that
    #   the first point goes to (1,0)
    #   the second point goes to (0,1)
    #   the third point goes to (0,0)
    grad0 = A.T @ gradp0
    grad1 = A.T @ gradp1
    grad2 = A.T @ gradp2
    allgrads = np.stack([grad0, grad1, grad2])
    
    # Fill in the matrices, considering pairs of basis functions
    # The minus in front of the Laplace operator will be added later.
    for ip1 in range(3):
        for ip2 in range(3):
            laplace[ip[ip1],ip[ip2]] += 0.5*detinvA*np.dot(allgrads[ip1,:], allgrads[ip2,:])
    
    for ip1 in range(3):
        for ip2 in range(3):
            bfint = 1/12 if ip1 == ip2 else 1/24
            wavenumop[ip[ip1],ip[ip2]] += k0**2*eps[ie]*detinvA * bfint
    
    for ip1 in range(3):
        for ip2 in range(3):
            bfint = 1/12 if ip1 == ip2 else 1/24
            mass[ip[ip1],ip[ip2]] += detinvA * bfint
t2 = time.time()
print(f'Construction: {t2-t1}')

A = -laplace + wavenumop
A = A.tocsc()
mass = mass.tocsc()
# Eigenvectors
t1 = time.time()
val,vec = eigs(A, M=mass, which='LR', k=5)
t2 = time.time()
print(f'Solution: {t2-t1}')
print('Mode indices:')
print(np.sqrt(val)/k0)

# Plot the solution
plt.rcParams['font.size'] = 18
# matplotlib wants its own triangulation object
mpltri = Triangulation(tri.points[:,0], tri.points[:,1], triangles=tri.simplices)
fig,ax = plt.subplots(1,5)
for i in range(5):
    ax[i].tripcolor(mpltri, abs(vec[:,i]), shading='gouraud')
    ax[i].set_title(f'{np.sqrt(val[i])/k0:.5}')
plt.show()

#plt.figure()
#plt.scatter(tri.points[:,0], tri.points[:,1], c=abs(vec[:,1]))
#plt.show()