"""Valley photonic crystal on silicon-on-insulator: transmission of three interfaces.

Running this file launches four simulations (a plain strip and the straight,
Z-shape and Omega-shape walls) on the cloud GPU. The HTML page is built from
recorded runs instead; its build never calls the run cell.

Device dimensions: He et al., Nat. Commun. 10, 872 (2019),
https://doi.org/10.1038/s41467-019-08881-z
"""

# %% imports
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.path import Path
import photonhub as ph
from photonhub.analysis import mode_launch, mode_monitor, solve_mode_on_cross_section

# %% parameters
# Design (µm): lattice constant, the two hole diameters, silicon thickness, strip width.
a, d1, d2, t, w = 0.385, 0.081, 0.181, 0.220, 1.700
n_si, n_sio2 = 3.47, 1.45
# Each interface is a chain of bearded cut lines: (heading in degrees, row index).
paths = {"strip": [],                 # no crystal: the straight strip the spectra are normalized to
         "flat": [(0, 0)],
         "Z": [(0, 0), (240, 14), (0, -6)],
         "Omega": [(0, 0), (120, -8), (0, 6), (240, 15), (0, 0)]}
n_len, n_rows = 23, 27                # crystal: lattice constants long, rows tall
facet_holes = True                    # keep the holes centred on the crystal's two end faces
band, n_lam = (1.320, 1.570), 126       # wavelength band (µm) and points
lam_map, show = 1.430, "Z"            # field map: Fig. 2e's wavelength and interface
# Simulation choices.
cells_x, cells_y = 19, 18             # cells per lattice constant and per row spacing
access, clear_xy, clear_z = 1.6, 0.5, 0.8
pml, record_ps = 16, 5.0
smoothing = "volume"                   # subpixel operator: volume-average permittivity in boundary cells

# %% materials
si = ph.Medium(permittivity=n_si ** 2)
air = ph.Medium(permittivity=1.0)
sio2 = ph.Medium(permittivity=n_sio2 ** 2)

# %% domain
h = a * np.sqrt(3) / 2                # row spacing
dx, dy = a / cells_x, h / cells_y
dz_si = t / (2 * int(np.ceil(t / (2 * min(dx, dy)))))
dz_max = min(band) / (n_sio2 * 20)
steps, dz = [], dz_si
while sum(steps) < clear_z:
    dz = min(1.25 * dz, dz_max)
    steps.append(dz)
steps += [dz_max] * pml               # outward from the silicon
z_steps = steps[::-1] + [dz_si] * round(t / dz_si) + steps
z_nodes = tuple(float(round(z, 9)) for z in np.r_[0.0, np.cumsum(z_steps)][:-1])
z_lo = sum(steps)
zc = z_lo + t / 2
nx = round((n_len * a + 2 * access) / dx) + 2 * pml
ny = round((n_rows * h + 2 * clear_xy) / dy) + 2 * pml
size = (nx * dx, ny * dy, sum(z_steps))
x0 = round((pml * dx + access) / dx) * dx
grid = ph.GradedMesh(dl_um=dx, coords=ph.GradedMeshAxis(
    y=tuple(round(j * dy, 9) for j in range(ny)), z=z_nodes))

# %% geometry
def crystal(lines):
    """Substrate, silicon and holes for one interface path; an empty path is the plain strip."""
    solid = [ph.Structure(medium=sio2, geometry=ph.Box(center_um=(size[0] / 2, size[1] / 2, 0.0),
                                                        size_um=(2 * size[0], 2 * size[1], 2 * z_lo)))]
    if not lines:
        strip = ph.Box(center_um=(size[0] / 2, size[1] / 2, zc), size_um=(size[0] + 2, w, t))
        return solid + [ph.Structure(medium=si, geometry=strip)], size[1] / 2, size[1] / 2
    normal = lambda deg: np.array([-np.sin(np.radians(deg)), np.cos(np.radians(deg))])
    offset = lambda m: (m - 1 / 6) * h
    corners = [np.linalg.solve([normal(p), normal(q)], [offset(i), offset(j)]) for (p, i), (q, j) in zip(lines, lines[1:])]
    ys = [offset(lines[0][1]), offset(lines[-1][1])] + [c[1] for c in corners]
    mid = (min(ys) + max(ys)) / 2
    y0, far = round((size[1] / 2 - mid) / dy) * dy, 10 * size[0]
    upper = Path(np.array([(-far, ys[0]), *corners, (far, ys[1]), (far, far), (-far, far)]) + [x0, y0])
    holes = []
    for m in range(-2 * n_rows, 2 * n_rows):
        for n in range(-2 * n_len, 3 * n_len):
            for y, big_above in ((m * h, True), (m * h + 2 * h / 3, False)):
                x = n * a + m * a / 2
                inside = -1e-9 < x < n_len * a + 1e-9 if facet_holes else 0 < x < n_len * a
                if inside and abs(y - mid) < n_rows * h / 2:
                    d = d2 if upper.contains_point((x0 + x, y0 + y)) == big_above else d1
                    holes.append(ph.Structure(medium=air, geometry=ph.Cylinder(
                        axis="z", center_um=(x0 + x, y0 + y, zc), radius_um=d / 2, length_um=t)))
    y_in, y_out = y0 + ys[0], y0 + ys[1]
    island = ph.Box(center_um=(x0 + n_len * a / 2, y0 + mid, zc), size_um=(n_len * a, n_rows * h, t))
    strips = [ph.Box(center_um=(x0 / 2, y_in, zc), size_um=(x0 + 2 * dx + a / 2, w, t)),
              ph.Box(center_um=((x0 + n_len * a + size[0]) / 2, y_out, zc), size_um=(size[0] - x0 - n_len * a + 2 * dx + a / 2, w, t))]
    return solid + [ph.Structure(medium=si, geometry=g) for g in [island] + strips] + holes, y_in, y_out

# %% band
c0 = 299792458.0
lam = np.linspace(*band, n_lam)
freqs = tuple(float(c0 / (value * 1e-6)) for value in lam)
f0 = 0.5 * (freqs[0] + freqs[-1])
f_map = c0 / (lam_map * 1e-6)
pulse = ph.GaussianPulse.for_band(wavelengths_um=list(band), freq0_hz=f0)
scene = dict(size_um=size, grid=grid, background=ph.Background(permittivity=1.0),
             boundaries=ph.Boundaries(x="pml", y="pml", z="pml"), pml_num_layers=pml,
             subpixel=True, subpixel_method=smoothing)
run = ph.RunSpec(run_time_s=record_ps * 1e-12, shutoff=1e-6)

# %% ports
window = dict(half_w_um=w / 2 + 0.8, half_v_um=t / 2 + 0.6, v_center_um=zc, dl_um=dx)
x_src, x_out = pml * dx + 0.4, size[0] - pml * dx - 0.4
sims, mon = {}, {}
for shape, lines in paths.items():
    structures, y_in, y_out = crystal(lines)
    basis = ph.Simulation(**scene, run=run, structures=tuple(structures))
    mode_in = solve_mode_on_cross_section(basis, "x", x_src, lam_map, "TE", 0, h_center_um=y_in, **window)
    mode_out = solve_mode_on_cross_section(basis, "x", x_out, lam_map, "TE", 0, h_center_um=y_out, **window)
    source = mode_launch(basis, mode_in, axis="x", position_um=x_src, source_time=pulse, direction="+")
    mon[shape] = {"in": mode_monitor(basis, mode_in, axis="x", position_um=x_src + 0.4, freqs_hz=freqs, name="in"),
                  "out": mode_monitor(basis, mode_out, axis="x", position_um=x_out, freqs_hz=freqs, name="out")}
    flux = ph.PowerMonitor(name="flux_out", axis="x", position_um=x_out, freqs_hz=freqs)
    field = ph.ProfileMonitor(name="field", center_um=(size[0] / 2, size[1] / 2, zc),
                              size_um=(size[0] - 2 * dx, size[1] - 2 * dy, 0.0),
                              fields=("Ex", "Ey", "Ez"), freqs_hz=(f_map,), interval_space=(2, 2, 1))
    planes = tuple(m.field_monitor for m in mon[shape].values()) + (flux,) + ((field,) if lines else ())
    sims[shape] = ph.Simulation(**scene, run=run, structures=tuple(structures),
                                sources=tuple(source), monitors=planes)

# %% simulation
sims[show].plot_index(z=zc)
plt.show()

# %% estimate
print(sims[show].cost_estimate().summary())

# %% run
data = {shape: ph.cloud.run_quoted(sim, max_usd=2.0) for shape, sim in sims.items()}

# %% field
titles = {"flat": "straight", "Z": "Z-shape, two bends", "Omega": "\u03a9-shape, four bends"}
walls = [s for s in data if s != "strip"]
fig, axes = plt.subplots(len(walls), 1, figsize=(6.6, 3.7 * len(walls)), constrained_layout=True, sharex=True)
for ax, shape in zip(np.atleast_1d(axes), walls):
    data[shape].plot_field("field", field="E", val="abs", z=zc, freq=f_map, simulation=sims[shape],
                           structures=False, scale="db", db_floor=-30, ax=ax)
    ax.annotate(titles[shape], xy=(0.02, 0.03), xycoords="axes fraction", color="white")
plt.show()

# %% readout
T_db = {}
for shape in paths:
    pin = mon[shape]["in"].mode_power(data[shape], direction="+")
    flux = data[shape]["flux_out"].sel(f=list(freqs)).values
    T_db[shape] = 10 * np.log10(flux / np.array([pin[f] for f in freqs]))
T_db = {shape: T_db[shape] - T_db["strip"] for shape in paths if paths[shape]}
ax = ph.viz.plot_comparison(lam * 1e3, T_db["flat"], label="straight", ylabel="transmission (dB)", ylim=(-30, 1))
for shape in [s for s in T_db if s != "flat"]:
    ax.plot(lam * 1e3, T_db[shape], label={"Z": "Z-shape, two bends", "Omega": "\u03a9-shape, four bends"}[shape])
ax.legend()
plt.show()
gap = (lam > 1.360) & (lam < 1.492)
print("mean in the band gap (dB):", {shape: round(float(v[gap].mean()), 2) for shape, v in T_db.items()})