# -*- coding: utf-8 -*-

import meep as mp
import math
import cmath
import numpy as np

resolution = 50        # pixels/μm

dpml = 1.0             # PML thickness
dsub = 3.0             # substrate thickness
dpad = 3.0             # length of padding between grating and PML
gp = 10.0              # grating period
gh = 0.5               # grating height
gdc = 0.5              # grating duty cycle

sx = dpml+dsub+gh+dpad+dpml
sy = gp

cell_size = mp.Vector3(sx, sy, 0)
pml_layers = [mp.PML(thickness=dpml, direction=mp.X)]

wvl = 0.5              # center wavelength
fcen = 1/wvl           # center frequency
df = 0.05*fcen         # frequency width

ng = 1.5
glass = mp.Medium(index=ng)

use_cw_solver = False  # CW solver or time stepping?
tol = 1e-6             # CW solver tolerance
max_iters = 2000       # CW solver max iterations
L = 10                 # CW solver L

# rotation angle of incident planewave; counter clockwise (CCW) about Z axis, 0 degrees along +X axis
theta_in = math.radians(10.7)

# k (in source medium) with correct length (plane of incidence: XY)
k = mp.Vector3(fcen*ng).rotate(mp.Vector3(z=1), theta_in)

symmetries = []
eig_parity = mp.ODD_Z
if theta_in == 0:
    k = mp.Vector3(0, 0, 0)
    symmetries = [mp.Mirror(mp.Y)]
    eig_parity += mp.EVEN_Y


def pw_amp(k, x0):
    def _pw_amp(x):
        return cmath.exp(1j*2*math.pi*k.dot(x+x0))
    return _pw_amp


src_pt = mp.Vector3(-0.5*sx+dpml+0.3*dsub, 0, 0)
sources = [mp.Source(mp.ContinuousSource(fcen, fwidth=df) if use_cw_solver else mp.GaussianSource(fcen, fwidth=df),
                     component=mp.Ez,
                     center=src_pt,
                     size=mp.Vector3(0, sy, 0),
                     amp_func=pw_amp(k, src_pt))]

sim = mp.Simulation(resolution=resolution,
                    cell_size=cell_size,
                    boundary_layers=pml_layers,
                    k_point=k,
                    default_material=glass,
                    sources=sources,
                    symmetries=symmetries)

refl_pt = mp.Vector3(-0.5*sx+dpml+0.5*dsub, 0, 0)
refl_flux = sim.add_flux(fcen, 0, 1, mp.FluxRegion(
    center=refl_pt, size=mp.Vector3(0, sy, 0)))

if use_cw_solver:
    sim.init_sim()
    sim.solve_cw(tol, max_iters, L)
else:
    sim.run(until_after_sources=100)

input_flux = mp.get_fluxes(refl_flux)
input_flux_data = sim.get_flux_data(refl_flux)

sim.reset_meep()

geometry = [mp.Block(material=glass, size=mp.Vector3(dpml+dsub, mp.inf, mp.inf), center=mp.Vector3(-0.5*sx+0.5*(dpml+dsub), 0, 0)),
            mp.Block(material=glass, size=mp.Vector3(gh, gdc*gp, mp.inf), center=mp.Vector3(-0.5*sx+dpml+dsub+0.5*gh, 0, 0))]

sim = mp.Simulation(resolution=resolution,
                    cell_size=cell_size,
                    boundary_layers=pml_layers,
                    geometry=geometry,
                    k_point=k,
                    sources=sources,
                    symmetries=symmetries)

refl_flux = sim.add_flux(fcen, 0, 1, mp.FluxRegion(
    center=refl_pt, size=mp.Vector3(0, sy, 0)))
sim.load_minus_flux_data(refl_flux, input_flux_data)

tran_pt = mp.Vector3(0.5*sx-dpml-0.5*dpad, 0, 0)
tran_flux = sim.add_flux(fcen, 0, 1, mp.FluxRegion(
    center=tran_pt, size=mp.Vector3(0, sy, 0)))

if use_cw_solver:
    sim.init_sim()
    sim.solve_cw(tol, max_iters, L)
else:
    sim.run(until_after_sources=200)

nm_r = np.floor((fcen*ng-k.y)*gp)-np.ceil((-fcen*ng-k.y)
                                          * gp)  # number of reflected orders
if theta_in == 0:
    nm_r = nm_r/2  # since eig_parity removes degeneracy in y-direction
nm_r = int(nm_r)

res = sim.get_eigenmode_coefficients(
    refl_flux, range(1, nm_r+1), eig_parity=eig_parity)
r_coeffs = res.alpha

Rsum = 0
for nm in range(nm_r):
    r_kdom = res.kdom[nm]
    Rmode = abs(r_coeffs[nm, 0, 1])**2/input_flux[0]
    r_angle = np.sign(r_kdom.y)*math.acos(r_kdom.x/(ng*fcen))
    print("refl:, {}, {:.2f}, {:.8f}".format(nm, math.degrees(r_angle), Rmode))
    Rsum += Rmode

# number of transmitted orders
nm_t = np.floor((fcen-k.y)*gp)-np.ceil((-fcen-k.y)*gp)
if theta_in == 0:
    nm_t = nm_t/2  # since eig_parity removes degeneracy in y-direction
nm_t = int(nm_t)

res = sim.get_eigenmode_coefficients(
    tran_flux, range(1, nm_t+1), eig_parity=eig_parity)
t_coeffs = res.alpha

Tsum = 0
for nm in range(nm_t):
    t_kdom = res.kdom[nm]
    Tmode = abs(t_coeffs[nm, 0, 0])**2/input_flux[0]
    t_angle = np.sign(t_kdom.y)*math.acos(t_kdom.x/fcen)
    print("tran:, {}, {:.2f}, {:.8f}".format(nm, math.degrees(t_angle), Tmode))
    Tsum += Tmode

print("mode-coeff:, {:.6f}, {:.6f}, {:.6f}".format(Rsum, Tsum, Rsum+Tsum))

r_flux = mp.get_fluxes(refl_flux)
t_flux = mp.get_fluxes(tran_flux)
Rflux = -r_flux[0]/input_flux[0]
Tflux = t_flux[0]/input_flux[0]
print(
    "poynting-flux:, {:.6f}, {:.6f}, {:.6f}".format(Rflux, Tflux, Rflux+Tflux))
