### compute the transmitted diffraction orders of a binary grating using mode decomposition
### based on two different methods: (1) MPB eigensolver and (2) DiffractedPlanewave object.
### Also, verify that the total power in all the orders is equivalent to the Poynting flux.

### for normal incidence, compute only positive diff. orders (total transmittance <= 0.50)
### for oblique incidence, compute ALL diff. orders (total transmittance <= 1.00)

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

def binary_grating_diffraction(gp, gh, gdc, theta):

  resolution = 50        # pixels/μm

  dpml = 1.0             # PML thickness
  dsub = 3.0             # substrate thickness
  dpad = 3.0             # length of padding between grating and PML

  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)

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

  eig_parity = mp.EVEN_Z

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

  symmetries = []
  if theta_in == 0:
    k = mp.Vector3()
    eig_parity += mp.ODD_Y
    symmetries = [mp.Mirror(direction=mp.Y,phase=-1)]

  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,0)
  sources = [mp.Source(mp.GaussianSource(fcen,fwidth=df),
                       component=mp.Hz,
                       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)

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

  sim.run(until_after_sources=50)

  input_flux = mp.get_fluxes(tran_mon)

  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)

  tran_mon = sim.add_mode_monitor(fcen, 0, 1,
                                  mp.FluxRegion(center=tran_pt, size=mp.Vector3(0,sy,0)))

  sim.run(until_after_sources=100)

  # number of (non-evanescent) 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
  nm_t = int(nm_t)+1

  bands = range(1,nm_t+1)

  if theta_in == 0:
    orders = range(0,nm_t)
  else:
    orders = range(int(np.ceil((-fcen-k.y)*gp)),int(np.floor((fcen-k.y)*gp))+1)

  eig_sum = 0
  dp_sum = 0

  for band,order in zip(bands,orders):
    res = sim.get_eigenmode_coefficients(tran_mon, [band], eig_parity=eig_parity)
    if res is not None:
      tran_eig = abs(res.alpha[0,0,0])**2/input_flux[0]
      if theta_in == 0:
        tran_eig = 0.5*tran_eig
    else:
      tran_eig = 0
    eig_sum += tran_eig

    res = sim.get_eigenmode_coefficients(tran_mon, mp.DiffractedPlanewave((0,order,0),mp.Vector3(0,1,0),0,1))
    if res is not None:
      tran_dp = abs(res.alpha[0,0,0])**2/input_flux[0]
      if (theta_in == 0) and (order == 0):
        tran_dp = 0.5*tran_dp
    else:
      tran_dp = 0
    dp_sum += tran_dp

    if theta_in == 0:
      err = abs(tran_eig-tran_dp)/tran_eig
      print("tran:, {:2d}, {:.8f}, {:2d}, {:.8f}, {:.8f}".format(band,tran_eig,order,tran_dp,err))
    else:
      print("tran:, {:2d}, {:.8f}, {:2d}, {:.8f}".format(band,tran_eig,order,tran_dp))

  flux = mp.get_fluxes(tran_mon)
  t_flux = flux[0]/input_flux[0]
  if (theta_in == 0):
    t_flux = 0.5*t_flux

  err = abs(dp_sum-t_flux)/t_flux
  print("flux:, {:.8f}, {:.8f}, {:.8f}, {:.8f}".format(eig_sum,
                                                       dp_sum,
                                                       t_flux,
                                                       err))

if __name__ == '__main__':
  binary_grating_diffraction(2.6,0.4,0.3,0)
  binary_grating_diffraction(3.7,0.6,0.4,13.5)
