wl = 1
freq0 = td.C_0 / 1
fwidth = 0.1 * freq0


def getSim(theta, permittivity=3.5**2, direction="z", min_steps_per_wvl=20):
    """Build a planar-scattering simulation for a slab and surrounding monitors."""

    plane_size = 1.3

    medium = td.Medium(permittivity=permittivity)

    # Frequencies to sample (single point for steady-state response).
    freqs = [freq0]

    # Convert incidence and polarization angles to radians.
    theta_rad = np.deg2rad(theta)

    # Select injection plane and Bloch axes depending on propagation direction.
    if direction == "z":
        center = (0, 0, 1)
        size = [td.inf, td.inf, 0]
        axis_1 = 0
        axis_2 = 1

    elif direction == "y":
        center = (0, 1, 0)
        size = [td.inf, 0, td.inf]
        axis_1 = 0
        axis_2 = 2

    elif direction == "x":
        center = (1, 0, 0)
        size = [0, td.inf, td.inf]
        axis_1 = 1
        axis_2 = 2

    # Define the incident plane wave illumination.
    planewave_0 = td.PlaneWave(
        size=size,
        center=center,
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
        direction="-",
        angle_theta=theta_rad,
    )

    # Bloch boundaries.
    bloch_1 = td.Boundary.bloch_from_source(
        source=planewave_0, domain_size=plane_size, axis=axis_1, medium=td.Medium(permittivity=1)
    )

    bloch_2 = td.Boundary.bloch_from_source(
        source=planewave_0, domain_size=plane_size, axis=axis_2, medium=td.Medium(permittivity=1)
    )

    # Create closed surface monitors used later for force calculations.
    surface_mons = surface_mon(center=(0, 0, 0), size=(1, 1, 1), freqs=freqs)

    if direction == "x":
        bspec = td.BoundarySpec(y=bloch_1, z=bloch_2, x=td.Boundary.absorber(num_layers=80))
        center_box = [-55.5, 500, 500]
        size_box = [111, 1222, 1222]
        sim_size = [4, plane_size, plane_size]
    elif direction == "y":
        bspec = td.BoundarySpec(x=bloch_1, z=bloch_2, y=td.Boundary.absorber(num_layers=80))
        center_box = [500, -55.5, 500]
        size_box = [1222, 111, 1222]
        sim_size = [plane_size, 4, plane_size]
    elif direction == "z":
        bspec = td.BoundarySpec(x=bloch_1, y=bloch_2, z=td.Boundary.absorber(num_layers=80))
        center_box = [500, 500, -55.5]
        size_box = [1222, 1222, 111]
        sim_size = [plane_size, plane_size, 4]

    # Slab structure occupying the high-index region.
    box_0 = td.Structure(
        geometry=td.Box(center=center_box, size=size_box), name="box_0", medium=medium
    )

    # Assemble the full simulation with sources, monitors, and structure.
    sim = td.Simulation(
        size=sim_size,
        boundary_spec=bspec,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=min_steps_per_wvl),
        run_time=5e-12,
        sources=[planewave_0],
        monitors=surface_mons,
        structures=[box_0],
    )

    return sim


sim = getSim(45)
_ = sim.plot(y=0)

plt.show()

sims_x = {}
sims_y = {}
sims_z = {}
thetas = [0, 10, 20, 30, 40, 50]
for theta in thetas:
    sims_x[f"{theta}"] = getSim(theta, direction="x")
    sims_y[f"{theta}"] = getSim(theta, direction="y")
    sims_z[f"{theta}"] = getSim(theta, direction="z")

batch_x = web.Batch(simulations=sims_x, folder_name="OpticalForces_x")
batch_data_thetas_x = batch_x.run(path_dir="batchThetas_x")

batch_y = web.Batch(simulations=sims_y, folder_name="OpticalForces_y")
batch_data_thetas_y = batch_y.run(path_dir="batchThetas_y")

batch_z = web.Batch(simulations=sims_z, folder_name="OpticalForces_z")
batch_data_thetas_z = batch_z.run(path_dir="batchThetas_z")

axes = ["x", "y", "z"]
components = ["Fx", "Fy", "Fz"]

# Pre-allocate containers for the force components keyed by face axis.
forces = {axis: {comp: [] for comp in components} for axis in axes}

for theta_key in batch_data_thetas_x.keys():
    print(theta_key)

    # Evaluate the stress tensor on each pair of monitors.
    evaluations = {
        "x": integrateStressTensor(batch_data_thetas_x[theta_key]),
        "y": integrateStressTensor(batch_data_thetas_y[theta_key]),
        "z": integrateStressTensor(batch_data_thetas_z[theta_key]),
    }

    # Append the results to the appropriate component lists.
    for axis, (fx_val, fy_val, fz_val) in evaluations.items():
        for comp, value in zip(components, (fx_val, fy_val, fz_val)):
            forces[axis][comp].append(np.array(value))

# Unpack back into the original variable names for downstream cells.
FxT_x, FyT_x, FzT_x = forces["x"]["Fx"], forces["x"]["Fy"], forces["x"]["Fz"]
FxT_y, FyT_y, FzT_y = forces["y"]["Fx"], forces["y"]["Fy"], forces["y"]["Fz"]
FxT_z, FyT_z, FzT_z = forces["z"]["Fx"], forces["z"]["Fy"], forces["z"]["Fz"]

def interface_force_eps(eps, I=1, A=1, theta_i=0):
    import numpy as np

    theta_i = np.deg2rad(theta_i)
    n1 = 1.0
    n2 = np.sqrt(eps)

    theta_t = np.arcsin(np.sin(theta_i) / n2)

    # Fresnel power coefficients (TE and TM)
    R = (
        (n2 * np.cos(theta_i) - n1 * np.cos(theta_t))
        / (n2 * np.cos(theta_i) + n1 * np.cos(theta_t))
    ) ** 2

    T = 1 - R  # lossless

    Fz = (I * A / td.C_0) * (np.cos(theta_i) * (1 + R) - n2 * np.cos(theta_t) * T)

    return np.abs(Fz)


# Adjusting the incident intensity to match the simulation results
I0 = (1 / 1.3) ** 2
Fz_a_theta = interface_force_eps(eps=3.5**2, theta_i=np.array(thetas), I=I0)

fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=False)
axes[0].plot(thetas, np.array([np.sum(i) for i in FxT_x]), label="Fx")
axes[0].plot(thetas, np.array([np.sum(i) for i in FyT_x]), label="Fy")
axes[0].plot(thetas, np.array([np.sum(i) for i in FzT_x]), label="Fz")
axes[0].plot(thetas, np.array(Fz_a_theta), "o", label="Analytic")
axes[0].set_title("x")
axes[0].legend()

axes[1].plot(thetas, np.array([np.sum(i) for i in FxT_y]), label="Fx")
axes[1].plot(thetas, np.array([np.sum(i) for i in FyT_y]), label="Fy")
axes[1].plot(thetas, np.array([np.sum(i) for i in FzT_y]), label="Fz")
axes[1].plot(thetas, np.array(Fz_a_theta), "o", label="Analytic")
axes[1].set_title("y")
axes[1].legend()

axes[2].plot(thetas, np.array([np.sum(i) for i in FxT_z]), label="Fx")
axes[2].plot(thetas, np.array([np.sum(i) for i in FyT_z]), label="Fy")
axes[2].plot(thetas, np.array([np.sum(i) for i in FzT_z]), label="Fz")
axes[2].plot(thetas, np.array(Fz_a_theta), "o", label="Analytic")
axes[2].set_title("z")
axes[2].legend()


for ax in axes:
    ax.set_xlabel("Angle (deg)")
    ax.set_ylabel("Force (N)")


plt.show()

from typing import List


def build_double_waveguide_sim(
    a: float = 0.310,
    d: float = 0.25,
    wavelength: float = 1.55,
    wg_medium: td.Medium | None = td.Medium(permittivity=3.45**2),
    x_span: float = 6.0,
    buffer_y: float = 2,
    buffer_z: float = 2,
    source_offset: float = 0.2,
    mode_index=1,
) -> td.Simulation:
    """Create a Tidy3D simulation containing two identical square waveguides"""

    if a <= 0.0:
        raise ValueError("Waveguide size 'a' must be positive.")
    if d < 0.0:
        raise ValueError("Waveguide separation 'd' must be non-negative.")
    if source_offset >= x_span / 2:
        raise ValueError("Source offset must be smaller than half of the simulation span.")

    # Distance from simulation centre to each waveguide centre along y.
    half_center_sep = 0.5 * (a + d)
    waveguide_centers: List[float] = [-half_center_sep, half_center_sep]

    # Simulation extents along each dimension (µm).
    y_half_span = a + d / 2 + buffer_y
    z_half_span = a / 2 + buffer_z
    sim_size = (x_span, 2 * y_half_span, 2 * z_half_span)

    structures = [
        td.Structure(
            geometry=td.Box(center=(0.0, yc, 0.0), size=(td.inf, a, a)),
            medium=wg_medium,
        )
        for yc in waveguide_centers
    ]

    freq0 = td.C_0 / wavelength
    pulse = td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10)
    source_x = -0.5 * x_span + source_offset

    sources = [
        td.ModeSource(
            center=(source_x, 0, 0.0),
            size=(0.0, td.inf, td.inf),
            source_time=pulse,
            direction="+",
            mode_spec=td.ModeSpec(num_modes=4),
            mode_index=mode_index,
        )
    ]

    sim = td.Simulation(
        size=sim_size,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=50),
        structures=structures,
        sources=sources,
        run_time=5e-12,
        monitors=surface_mon(
            center=(0, (d / 2 + a / 2 + 0.5), 0), size=(1, a + d / 2 + 1, 1), freqs=[freq0]
        ),
        symmetry=(0, 0, -1),
    )

    ms = ModeSolver(
        simulation=sim,
        plane=sources[0].bounding_box,
        freqs=[sources[0].source_time.freq0],
        mode_spec=sources[0].mode_spec,
    )

    return sim, ms


sim, ms = build_double_waveguide_sim(d=0.1)

mode_data = web.run(ms, "mode")

fig, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)

for ax, (mode_index, title) in zip(
    axes,
    [
        (0, "E_z (symmetric, mode 0)"),
        (1, "E_z (antisymmetric, mode 1)"),
    ],
):
    ms.plot_field("Ez", "real", mode_index=mode_index, ax=ax)
    ax.set_title(title)

# 0 → symmetric fundamental mode
# 1 → antisymmetric higher-order mode

plt.show()

d_values = np.linspace(0.03, 0.31, 21)
sims_1 = {str(i): build_double_waveguide_sim(d=i, mode_index=0)[0] for i in d_values}
batch_1 = web.Batch(simulations=sims_1, folder_name="double_waveguide_1")
batch_data_1 = batch_1.run(path_dir="double_waveguide_1")

sims_2 = {str(i): build_double_waveguide_sim(d=i, mode_index=1)[0] for i in d_values}
batch_2 = web.Batch(simulations=sims_2, folder_name="double_waveguide_2")
batch_data_2 = batch_2.run(path_dir="double_waveguide_2")

F_1 = [integrateStressTensor(batch_data_1[i]) for i in batch_data_1.keys()]
F_2 = [integrateStressTensor(batch_data_2[i]) for i in batch_data_2.keys()]

fig, ax = plt.subplots()
factor = 0.31 * td.C_0 / 1  # Source is normalized to 1W
ax.plot(d_values / 0.31, np.array([i[1] for i in F_2]) * factor, label="Antisymmetric")
ax.plot(d_values / 0.31, np.array([i[1] for i in F_1]) * factor, label="Symmetric")
ax.set_xlabel("d/a")
ax.set_ylabel("Force (N/L)(ac/P)")

ax.legend()
plt.show()