from mpl_toolkits.axes_grid1 import make_axes_locatable
import meep as mp
import numpy as np
import math
import matplotlib.pyplot as plt

# True:  plot the scattered fields in the extended air region adjacent to the grating
# False: plot the diffraction spectra based on a 1d cross section of the scattered fields
field_profile = True

resolution = 50         # pixels/μm

dpml = 1.0              # PML thickness
dsub = 2.0              # substrate thickness
dpad = 1.0              # flat-surface padding
gp = 1.0                # grating periodicity
gh = 0.5                # grating height
gdc = 0.5               # grating duty cycle
num_cells = 5           # number of grating unit cells

# air region thickness adjacent to grating
dair = 10 if field_profile else dpad

wvl = 0.5               # center wavelength
fcen = 1/wvl            # center frequency

k_point = mp.Vector3()

glass = mp.Medium(index=1.5)

pml_layers = [mp.PML(thickness=dpml)]

symmetries = [mp.Mirror(mp.Y)]

sx = dpml+dsub+gh+dair+dpml
sy = dpml+dpad+num_cells*gp+dpad+dpml
cell_size = mp.Vector3(sx, sy)

src_pt = mp.Vector3(-0.5*sx+dpml+0.5*dsub)
sources = [mp.Source(mp.GaussianSource(fcen, fwidth=0.2*fcen, is_integrated=True),
                     component=mp.Ez,
                     center=src_pt,
                     size=mp.Vector3(y=sy))]

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

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

mon_pt = mp.Vector3(0.5*sx-dpml-0.5*dair)
near_fields = sim.add_dft_fields([mp.Ez], fcen, fcen, 1, center=mon_pt, size=mp.Vector3(
    dair if field_profile else 0, sy-2*dpml))

sim.run(until_after_sources=100)

flat_dft = sim.get_dft_array(near_fields, mp.Ez, 0)

sim.reset_meep()

for j in range(num_cells):
    geometry.append(mp.Block(material=glass,
                             size=mp.Vector3(gh, gdc*gp, mp.inf),
                             center=mp.Vector3(-0.5*sx+dpml+dsub+0.5*gh, -0.5*sy+dpml+dpad+(j+0.5)*gp)))

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

near_fields = sim.add_dft_fields([mp.Ez], fcen, fcen, 1, center=mon_pt, size=mp.Vector3(
    dair if field_profile else 0, sy-2*dpml))

sim.run(until_after_sources=100)

grating_dft = sim.get_dft_array(near_fields, mp.Ez, 0)

scattered_field = grating_dft-flat_dft
scattered_amplitude = np.abs(scattered_field)**2

[x, y, z, w] = sim.get_array_metadata(dft_cell=near_fields)

if field_profile:
    if mp.am_master():
        plt.figure(dpi=150)
        plt.pcolormesh(x, y, np.rot90(scattered_amplitude), cmap='inferno',
                                                                 shading='gouraud', vmin=0, vmax=scattered_amplitude.max())
        plt.gca().set_aspect('equal')
        plt.xlabel('x (μm)')
        plt.ylabel('y (μm)')

        # ensure that the height of the colobar matches that of the plot
        from mpl_toolkits.axes_grid1 import make_axes_locatable
        divider = make_axes_locatable(plt.gca())
        cax = divider.append_axes("right", size="5%", pad=0.05)
        plt.colorbar(cax=cax)
        plt.tight_layout()
        plt.show()
else:
    ky = np.fft.fftshift(np.fft.fftfreq(len(scattered_field), 1/resolution))
    FT_scattered_field = np.fft.fftshift(np.fft.fft(scattered_field))
    if mp.am_master():
        plt.figure(dpi=150)
        plt.subplots_adjust(hspace=0.3)

        plt.subplot(2, 1, 1)
        plt.plot(y, scattered_amplitude, 'bo-')
        plt.xlabel("y (μm)")
        plt.ylabel("field amplitude")

        plt.subplot(2, 1, 2)
        plt.plot(ky, np.abs(FT_scattered_field)**2, 'ro-')
        plt.gca().ticklabel_format(axis='y', style='sci', scilimits=(0, 0))
        plt.xlabel(r'wavevector k$_y$, 2π (μm)$^{-1}$')
        plt.ylabel("Fourier transform")
        plt.gca().set_xlim([-3, 3])

        plt.tight_layout(pad=1.0)
        plt.show()

plt.figure(dpi=150)
plt.pcolormesh(x, y, np.rot90(scattered_amplitude), cmap='inferno',
                                                         shading='gouraud', vmin=0, vmax=scattered_amplitude.max())
plt.gca().set_aspect('equal')
plt.xlabel('x (μm)')
plt.ylabel('y (μm)')
# ensure that the height of the colobar matches that of the plot
divider = make_axes_locatable(plt.gca())
cax = divider.append_axes("right", size="5%", pad=0.05)
plt.colorbar(cax=cax)
plt.tight_layout()
plt.show()
