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

import meep as mp
import numpy as np
import matplotlib.pyplot as plt
import numpy.matlib
import argparse

resolution = 60         # pixels/μm

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

wvl_min = 0.4           # min wavelength
wvl_max = 0.6           # max wavelength
fmin = 1/wvl_max        # min frequency
fmax = 1/wvl_min        # max frequency
fcen = 0.5*(fmin+fmax)  # center frequency
df = fmax-fmin          # frequency width
nfreq = 21              # number of frequency bins

k_point = mp.Vector3(0, 0, 0)

glass = mp.Medium(index=1.5)


def grating(gp, gh, gdc, oddz):
    sx = dpml+dsub+gh+dpad+dpml
    sy = gp

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

    src_pt = mp.Vector3(-0.5*sx+dpml+0.5*dsub, 0, 0)
    sources = [mp.Source(mp.GaussianSource(
        fcen, fwidth=df), component=mp.Ez if oddz else mp.Hz, center=src_pt, size=mp.Vector3(0, sy, 0))]

    symmetries = [mp.Mirror(mp.Y, phase=+1 if oddz else -1)]

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

    mon_pt = mp.Vector3(0.5*sx-dpml-0.5*dpad, 0, 0)
    flux_mon = sim.add_flux(fcen, df, nfreq, mp.FluxRegion(
        center=mon_pt, size=mp.Vector3(0, sy, 0)))

    sim.run(until_after_sources=100)

    input_flux = mp.get_fluxes(flux_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_point,
                        sources=sources,
                        symmetries=symmetries)

    mode_mon = sim.add_flux(fcen, df, nfreq, mp.FluxRegion(
        center=mon_pt, size=mp.Vector3(0, sy, 0)))

    sim.run(until_after_sources=300)

    freqs = mp.get_eigenmode_freqs(mode_mon)
    res = sim.get_eigenmode_coefficients(
        mode_mon, [1], eig_parity=mp.ODD_Z+mp.EVEN_Y if oddz else mp.EVEN_Z+mp.ODD_Y)
    coeffs = res.alpha

    mode_wvl = [1/freqs[nf] for nf in range(nfreq)]
    mode_tran = [abs(coeffs[0, nf, 0])**2/input_flux[nf]
                 for nf in range(nfreq)]
    mode_phase = [np.angle(coeffs[0, nf, 0]) for nf in range(nfreq)]

    return mode_wvl, mode_tran, mode_phase


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-gp', type=float, default=0.35,
                        help='grating periodicity (default: 0.35 μm)')
    parser.add_argument('-gh', type=float, default=0.6,
                        help='grating height (default: 0.6 μm)')
    parser.add_argument('-oddz', action='store_true',
                        default=False, help='oddz? (default: False)')
    args = parser.parse_args()

    gdc = np.arange(0.1, 1.0, 0.1)
    mode_tran = np.empty((gdc.size, nfreq))
    mode_phase = np.empty((gdc.size, nfreq))
    for n in range(gdc.size):
        mode_wvl, mode_tran[n, :], mode_phase[n, :] = grating(
            args.gp, args.gh, gdc[n], args.oddz)

    plt.figure(dpi=150)

    plt.subplot(1, 2, 1)
    plt.pcolormesh(mode_wvl, gdc, mode_tran, cmap='hot_r',
                   shading='gouraud', vmin=0, vmax=mode_tran.max())
    plt.axis([wvl_min, wvl_max, gdc[0], gdc[-1]])
    plt.xlabel("wavelength (μm)")
    plt.xticks([t for t in np.arange(wvl_min, wvl_max+0.1, 0.1)])
    plt.ylabel("grating duty cycle")
    plt.yticks([t for t in np.arange(gdc[0], gdc[-1]+0.1, 0.1)])
    plt.title("transmittance")
    cbar = plt.colorbar()
    cbar.set_ticks([t for t in np.arange(0, 1.2, 0.2)])
    cbar.set_ticklabels(["{:.1f}".format(t) for t in np.arange(0, 1.2, 0.2)])

    plt.subplot(1, 2, 2)
    plt.pcolormesh(mode_wvl, gdc, mode_phase, cmap='RdBu',
                   shading='gouraud', vmin=mode_phase.min(), vmax=mode_phase.max())
    plt.axis([wvl_min, wvl_max, gdc[0], gdc[-1]])
    plt.xlabel("wavelength (μm)")
    plt.xticks([t for t in np.arange(wvl_min, wvl_max+0.1, 0.1)])
    plt.ylabel("grating duty cycle")
    plt.yticks([t for t in np.arange(gdc[0], gdc[-1]+0.1, 0.1)])
    plt.title("phase (radians)")
    cbar = plt.colorbar()
    cbar.set_ticks([t for t in range(-3, 4)])
    cbar.set_ticklabels(["{:.1f}".format(t) for t in range(-3, 4)])

    plt.tight_layout()
    plt.show()
