# standard python imports
import matplotlib.pyplot as plt
import numpy as np
import tidy3d as td
from tidy3d import web

td.config.logging.level = "ERROR"
td.config.local_cache.enabled = False

# The setup of the 1D simulation:8

# |     |          |        |            |         |            |     |
# |`PML`|   dPML   |`source`|  distance  |`monitor`|  distance  |`PML`|
# |     |          |        |            |         |            |     |

# Number of PML layers
Npml = 6  # the default number of layers is 12

# sufficiently large distance between the source and the time monitor
distance = 100  # um
# distance between the source and PML
dPML = 1  # um

# Wavelength and frequency range of the source
freq0 = 300e12  # Hz, central frequency
fwidth = freq0 / 5  # frequency pulse width
lambda0 = td.C_0 / freq0

# resolution
resolution = 20
dl = lambda0 / resolution

# runtime
t_stop = 1e-11

# simulation domain size
sim_size = Lx, Ly, Lz = (dl, dl, 2 * distance + dPML)

# Uniform grid
grid_spec = td.GridSpec.uniform(dl=dl)

# Planewave source
src_z = dPML - Lz / 2
source = td.PlaneWave(
    center=(0, 0, src_z),
    size=(td.inf, td.inf, 0),
    source_time=td.GaussianPulse(
        freq0=freq0,
        fwidth=fwidth,
    ),
    direction="+",
    name="Gaussian pulse",
)

# Time monitor at a single point
tmnt_pt = td.FieldTimeMonitor(
    center=(0, 0, src_z + distance),
    name="time",
    size=(0, 0, 0),
    interval=1,
)

# PML boundary
pml = td.Boundary.pml(num_layers=Npml)
bspec = td.BoundarySpec(x=td.Boundary.periodic(), y=td.Boundary.periodic(), z=pml)

sim = td.Simulation(
    center=(0, 0, 0),
    size=sim_size,
    grid_spec=grid_spec,
    structures=[],
    sources=[source],
    monitors=[tmnt_pt],
    run_time=t_stop,
    boundary_spec=bspec,
    shutoff=0,
)

sim_data = web.run(
    sim, task_name="lecture06_pml_reflection", path="data/data_pml_reflection.hdf5"
)

plt.rcParams.update({"font.size": 16})
fig, ax = plt.subplots(1, figsize=(7, 4), tight_layout=True)

speed = td.C_0
time_reach_monitor = source.source_time.offset / fwidth / 2 / np.pi + distance / speed
time_reflection = time_reach_monitor + 2 * distance / speed
ax.vlines(time_reach_monitor * 1e12, -600, 600, colors="r", linestyle="dashed")
ax.vlines(time_reflection * 1e12, -600, 600, colors="b", linestyle="dashed")

time_data = sim_data["time"]
Ex = time_data.Ex
Ex_values = np.squeeze(Ex.values)

ax.plot(Ex.t * 1e12, Ex_values)
ax.set_xlabel("time (ps)")
ax.set_ylabel("$E_x$")
ax.set_xlim(0, 1.5)
ax.set_ylim(-600, 600)
plt.show()

# First, let's select two time intervals of equal length. The first interval
# completely contains the original incident pulse, and the second interval
# completely contains the reflected pulse. We take the former to
# be [0,time_reach_PML), and the latter to be [time_reach_PML,2*time_reach_PML)

time_reach_PML = source.source_time.offset / fwidth / 2 / np.pi + 2 * distance / speed
inds = np.argwhere(np.array(Ex.t) >= time_reach_PML)
Nt = int(inds[0].item())

# Now let's perform FFT in each interval, and visualize the spectrum
dt = sim.dt
df = 1 / Nt / dt
fmesh = np.arange(0, Nt) * df

# original pulse
spectrum_original = (
    1 / np.sqrt(2 * np.pi) * np.fft.ifft(np.fft.ifftshift(Ex_values[:Nt])) / df
)
# reflected pulse
spectrum_reflection = (
    1 / np.sqrt(2 * np.pi) * np.fft.ifft(np.fft.ifftshift(Ex_values[Nt : Nt * 2])) / df
)

# visualization of each spectrum
fig, ax = plt.subplots(1, 2, figsize=(8, 4), tight_layout=True)
ax[0].plot(fmesh / 1e12, np.abs(spectrum_original), label="Original")
ax[0].plot(fmesh / 1e12, np.abs(spectrum_reflection), label="Reflected")
ax[0].set_xlim((freq0 - fwidth) / 1e12, (freq0 + fwidth) / 1e12)
ax[0].set_ylim(0, 8e-13)
ax[0].set_xlabel("Frequency (THz)")
ax[0].set_title("Spectrum")
ax[0].legend()

# visualization of reflection spectrum
ax[1].plot(fmesh / 1e12, np.abs(spectrum_reflection / spectrum_original))
ax[1].set_xlim((freq0 - fwidth) / 1e12, (freq0 + fwidth) / 1e12)
ax[1].set_yscale("log")
ax[1].set_ylim(1e-6, 1)
ax[1].set_title("Reflection amplitude")
ax[1].set_xlabel("Frequency (THz)")
plt.show()

# Wavelength and frequency range
freq_range = (200e12, 400e12)
lambda_range = (td.constants.C_0 / freq_range[1], td.constants.C_0 / freq_range[0])
freq0 = np.sum(freq_range) / 2

# frequencies and wavelengths of monitor
Nfreq = 301
monitor_freqs = np.linspace(freq_range[0], freq_range[1], Nfreq)
monitor_lambdas = td.constants.C_0 / monitor_freqs

# central frequency, frequency pulse width and total running time
lambda0 = td.C_0 / freq0
bandwidth = 0.4  # bandwidth of source in units of delta frequency.  0.38 for broadband

freqw = bandwidth * (freq_range[1] - freq_range[0])
t_stop = 1000 / freq0

# Thickness and refractive index of slab
t_slab = 0.5
n_slab = 3.5
mat_slab = td.Medium(permittivity=n_slab**2, name="silicon")

# Grid resolution (cells per um)
dl = 6e-3

# space between slabs and sources and PML
spacing = 16 * lambda_range[-1]

# simulation size
sim_size = Lx, Ly, Lz = (dl, dl, 2 * spacing + t_slab)