%matplotlib notebook
import netgen.gui
from netgen.geom2d import SplineGeometry
from ngsolve import *
import matplotlib.pyplot as plt

geo = SplineGeometry()
geo.AddCircle( (0,0), 1.7, leftdomain=3, bc="outerbnd")
geo.AddCircle( (0,0), 1.25, leftdomain=2, rightdomain=3, bc="midbnd")
geo.AddCircle( (0,0), 1, leftdomain=1, rightdomain=2, bc="innerbnd")
geo.SetMaterial(1, "inner")
geo.SetMaterial(2, "mid")
geo.SetMaterial(3, "pmlregion")
mesh = Mesh(geo.GenerateMesh (maxh=0.075))
mesh.Curve(3)

mesh.SetPML(pml.Radial(rad=1.25,alpha=5j,origin=(0,0)), "pmlregion") #Alpha is the strenth of the PML.

import numpy as np

omega_0 = 20
omega_tilde = 350 #70*exp(-(((x*x)*np.log(7/2))+((y*y)*np.log(7/2)))) #Gaussian function for our test Omega.

domain_values = {'inner': omega_tilde, 'mid': omega_0, 'pmlregion': omega_0}
values_list = [domain_values[mat] for mat in mesh.GetMaterials()]
omega = CoefficientFunction(values_list)

Draw(omega, mesh, 'piecewise')

fes = H1(mesh, complex=True, order=5)

s = np.pi

u_in =exp(1j*omega_0*(cos(s)*x + sin(s)*y)) #Can use any vector as long as it is on the unit circle. 

#Defining our test and solution functions.
u = fes.TrialFunction()
v = fes.TestFunction()

#Defining our LHS of the problem.
a = BilinearForm(fes)
a += grad(u)*grad(v)*dx - omega**2*u*v*dx
a += -omega*1j*u*v * ds("innerbnd")
a.Assemble()

#Defining the RHS of our problem.
f = LinearForm(fes)
f += u_in * (omega**2 - omega_0**2) * v * dx
f.Assemble()

#Solving our problem.
u_s = GridFunction(fes, name="u")
u_s.vec.data = a.mat.Inverse() * f.vec

u_tot = u_in + u_s

def f1_field(r,s):
    
    return Integrate(exp(-1j*omega_0*((cos(r)-cos(s))*x+(sin(r)-sin(s))*y))*omega_tilde,
              mesh,definedon=mesh.Materials("inner"))

r = 3/2 * np.pi

us_n = BoundaryFromVolumeCF(u_s)
ecomp = exp(-1j*omega_0*(cos(r)*x + sin(r)*y))
ecomp_n = BoundaryFromVolumeCF(ecomp)

def ff_field(r):
    
    temp = Integrate(u_s*ecomp_n - us_n*ecomp, mesh,definedon=mesh.Boundaries("innerbnd"))
    
    return (exp(1j*np.pi/4)/(8*np.pi*omega_0)**.5)*temp
    
    
point = mesh(1.7*cos(r),1.7*sin(r))

print("Full Comp:", ff_field(r), "\nFirst Order:", f1_field(r,s))