# Central wavelength and frequency used for all epsilon queries and monitors
wavelength = 1.55
freq0 = td.C_0 / wavelength

# Background medium
cladding_medium = td.Medium(permittivity=1.44**2)

# Layer media defined by refractive index
bottom_layer_medium = td.Medium(permittivity=2.0**2)
middle_layer_medium = td.Medium(permittivity=3.0**2)
top_layer_medium = td.Medium(permittivity=4.0**2)

# Lateral dimensions in micrometers
waveguide_width = 0.50
waveguide_length = 1e3

# Vertical stack definition in micrometers
layer_thickness = 0.10
stack_thickness = 3 * layer_thickness

# Simulation region dimensions in micrometers
sim_size_x = 4.0
sim_size_y = 2.0
sim_size_z = 4.0

# Mesh target
min_steps_per_wvl = 10

# Three stacked layers along z
bottom_layer_box = td.Box(
    center=(0.0, 0.0, -layer_thickness),
    size=(waveguide_length, waveguide_width, layer_thickness),
)

middle_layer_box = td.Box(
    center=(0.0, 0.0, 0.0),
    size=(waveguide_length, waveguide_width, layer_thickness),
)

top_layer_box = td.Box(
    center=(0.0, 0.0, layer_thickness),
    size=(waveguide_length, waveguide_width, layer_thickness),
)

bottom_layer = td.Structure(
    geometry=bottom_layer_box,
    medium=bottom_layer_medium,
    name="bottom_layer_n2",
)

middle_layer = td.Structure(
    geometry=middle_layer_box,
    medium=middle_layer_medium,
    name="middle_layer_n3",
)

top_layer = td.Structure(
    geometry=top_layer_box,
    medium=top_layer_medium,
    name="top_layer_n4",
)

stack_structures = [bottom_layer, middle_layer, top_layer]

# Box used for the local epsilon extraction.
# It is a line along z located at x=0 and y=0, crossing the stack vertically.
line_box = td.Box(
    center=(0.0, 0.0, 0.0),
    size=(0.0, 0.0, 1.0),
)

layer_boundaries = np.array(
    [
        -stack_thickness / 2,
        -stack_thickness / 2 + layer_thickness,
        -stack_thickness / 2 + 2 * layer_thickness,
        stack_thickness / 2,
    ]
)

def make_simulation(include_source_and_monitor=False):
    monitors = []
    sources = []

    if include_source_and_monitor:
        source_time = td.GaussianPulse(
            freq0=freq0,
            fwidth=0.2 * freq0,
        )

        plane_wave = td.PlaneWave(
            center=(0.0, 0.0, -0.6),
            size=(sim_size_x, sim_size_y, 0.0),
            source_time=source_time,
            direction="+",
        )

        permittivity_monitor = td.PermittivityMonitor(
            center=line_box.center,
            size=line_box.size,
            freqs=[freq0],
            name="eps_line_monitor",
        )

        sources = [plane_wave]
        monitors = [permittivity_monitor]

    sim = td.Simulation(
        size=(sim_size_x, sim_size_y, sim_size_z),
        center=(0.0, 0.0, 0.0),
        medium=cladding_medium,
        structures=stack_structures,
        sources=sources,
        monitors=monitors,
        grid_spec=td.GridSpec.auto(
            wavelength=wavelength,
            min_steps_per_wvl=min_steps_per_wvl,
        ),
        run_time=1e-14,
        boundary_spec=td.BoundarySpec.all_sides(td.PML()),
    )
    return sim

# Local epsilon without local subpixel averaging
td.config.simulation.use_local_subpixel = False
sim_no_subpixel = make_simulation(include_source_and_monitor=False)

eps_line_no_subpixel = sim_no_subpixel.epsilon(
    box=line_box,
    coord_key="Ex",
    freq=freq0,
)

# Local epsilon with local subpixel averaging
td.config.simulation.use_local_subpixel = True

sim_with_subpixel = make_simulation(include_source_and_monitor=False)

eps_line_with_subpixel = sim_with_subpixel.epsilon(
    box=line_box,
    coord_key="Ex",
    freq=freq0,
)

def extract_line_profile(eps_data):
    # For a line box, keep the non-singleton spatial axis.
    # The returned object can carry tensor components such as eps_xx, eps_yy, eps_zz.
    if hasattr(eps_data, "eps_xx"):
        data = eps_data.eps_xx
    else:
        data = eps_data

    squeezed = data.squeeze(drop=True)

    for axis_name in ("x", "y", "z"):
        if axis_name in squeezed.coords and squeezed.coords[axis_name].size > 1:
            coord = squeezed.coords[axis_name].values
            values = np.real(squeezed.values)
            return axis_name, coord, values

    raise RuntimeError("Could not identify the varying spatial coordinate in the epsilon data.")


axis_no_subpixel, coord_no_subpixel, values_no_subpixel = extract_line_profile(eps_line_no_subpixel)

axis_with_subpixel, coord_with_subpixel, values_with_subpixel = extract_line_profile(
    eps_line_with_subpixel
)

fig, ax = plt.subplots(figsize=(8, 5))

ax.plot(coord_no_subpixel, values_no_subpixel, "o-", label="epsilon(box), no subpixel")

ax.plot(
    coord_with_subpixel,
    values_with_subpixel,
    "o-",
    label="epsilon(box), with subpixel",
)

for boundary_index, boundary_position in enumerate(layer_boundaries):
    label = "layer boundary" if boundary_index == 0 else None
    ax.axvline(boundary_position, linestyle=":", alpha=0.5, label=label)

ax.set_xlabel(f"{axis_no_subpixel} (µm)")
ax.set_ylabel("Local permittivity")
ax.set_title("Local permittivity across the three-layer stack")
ax.legend()
ax.set_xlim(-0.2, 0.2)
plt.show()

sim_monitor = make_simulation(include_source_and_monitor=True)

sim_data = web.run(
    sim_monitor,
    task_name="subpixel_local_epsilon_dummy_monitor",
)

# Grab the monitored permittivity directly along the z line.
monitor_eps_xx = sim_data["eps_line_monitor"].eps_xx.squeeze(drop=True)
monitor_coord = monitor_eps_xx.coords["z"].values
monitor_values = np.real(monitor_eps_xx.values)

fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(
    coord_with_subpixel,
    values_with_subpixel,
    "o-",
    label="local epsilon with subpixel",
)
ax.plot(monitor_coord, monitor_values, "s--", label="PermittivityMonitor")

for boundary_index, boundary_position in enumerate(layer_boundaries):
    label = "layer boundary" if boundary_index == 0 else None
    ax.axvline(boundary_position, linestyle=":", alpha=0.5, label=label)

ax.set_xlabel("z (µm)")
ax.set_ylabel("Permittivity")
ax.set_title("Local epsilon vs PermittivityMonitor for the three-layer stack")
ax.legend()
plt.show()