"""
Waveguide Bend Flux Analysis with OptixLog Integration

Computes transmission and reflection through a 90-degree waveguide bend.
"""

import os
import meep as mp
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import numpy as np
from optixlog import Optixlog


def main():
    api_key = os.getenv("OPTIX_API_KEY", "your_api_key_here")
    
    client = Optixlog(api_key=api_key)
    project = client.project(name="MeepExamples", create_if_not_exists=True)

    # Simulation parameters
    resolution = 10
    sx = 16
    sy = 32
    dpml = 1.0
    pad = 4
    w = 1
    epsilon_wg = 12.0
    fcen = 0.15
    df = 0.1
    nfreq = 100

    wvg_xcen = 0.5 * (sx - w - 2 * pad)
    wvg_ycen = -0.5 * (sy - w - 2 * pad)

    run = project.run(
        name="waveguide_bend_flux",
        config={
            "simulation_type": "waveguide_bend_flux",
            "resolution": resolution,
            "waveguide_width": w,
            "waveguide_epsilon": epsilon_wg,
            "center_frequency": fcen,
            "num_frequencies": nfreq
        }
    )
    
    run.log(step=0, resolution=resolution, waveguide_width=w, center_frequency=fcen)

    cell = mp.Vector3(sx, sy, 0)
    pml_layers = [mp.PML(dpml)]

    # Straight waveguide geometry
    straight_geometry = [
        mp.Block(
            size=mp.Vector3(mp.inf, w, mp.inf),
            center=mp.Vector3(0, wvg_ycen, 0),
            material=mp.Medium(epsilon=epsilon_wg),
        )
    ]

    sources = [
        mp.Source(
            mp.GaussianSource(fcen, fwidth=df),
            component=mp.Ez,
            center=mp.Vector3(-0.5 * sx + dpml, wvg_ycen, 0),
            size=mp.Vector3(0, w, 0),
        )
    ]

    # ===== Straight waveguide simulation =====
    print("⚡ Running straight waveguide simulation...")
    sim = mp.Simulation(
        cell_size=cell,
        boundary_layers=pml_layers,
        geometry=straight_geometry,
        sources=sources,
        resolution=resolution,
    )

    refl_fr = mp.FluxRegion(
        center=mp.Vector3(-0.5 * sx + dpml + 0.5, wvg_ycen, 0),
        size=mp.Vector3(0, 2 * w, 0)
    )
    refl = sim.add_flux(fcen, df, nfreq, refl_fr)

    tran_fr = mp.FluxRegion(
        center=mp.Vector3(0.5 * sx - dpml, wvg_ycen, 0),
        size=mp.Vector3(0, 2 * w, 0)
    )
    tran = sim.add_flux(fcen, df, nfreq, tran_fr)

    pt = mp.Vector3(0.5 * sx - dpml - 0.5, wvg_ycen)
    sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, pt, 1e-3))

    # Save reference data
    straight_refl_data = sim.get_flux_data(refl)
    straight_tran_flux = mp.get_fluxes(tran)

    run.log(step=1, phase="straight_waveguide",
            max_transmission_flux=float(max(straight_tran_flux)))

    sim.reset_meep()

    # ===== Bent waveguide simulation =====
    print("⚡ Running bent waveguide simulation...")
    bent_geometry = [
        mp.Block(
            mp.Vector3(sx - pad, w, mp.inf),
            center=mp.Vector3(-0.5 * pad, wvg_ycen),
            material=mp.Medium(epsilon=epsilon_wg),
        ),
        mp.Block(
            mp.Vector3(w, sy - pad, mp.inf),
            center=mp.Vector3(wvg_xcen, 0.5 * pad),
            material=mp.Medium(epsilon=epsilon_wg),
        ),
    ]

    sim = mp.Simulation(
        cell_size=cell,
        boundary_layers=pml_layers,
        geometry=bent_geometry,
        sources=sources,
        resolution=resolution,
    )

    refl = sim.add_flux(fcen, df, nfreq, refl_fr)
    tran_fr = mp.FluxRegion(
        center=mp.Vector3(wvg_xcen, 0.5 * sy - dpml - 0.5, 0),
        size=mp.Vector3(2 * w, 0, 0)
    )
    tran = sim.add_flux(fcen, df, nfreq, tran_fr)

    # Subtract incident flux from reflection monitor
    sim.load_minus_flux_data(refl, straight_refl_data)

    pt = mp.Vector3(wvg_xcen, 0.5 * sy - dpml - 0.5)
    sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, pt, 1e-3))

    bend_refl_flux = mp.get_fluxes(refl)
    bend_tran_flux = mp.get_fluxes(tran)
    flux_freqs = mp.get_flux_freqs(refl)

    # Calculate transmission coefficients
    wavelengths = []
    reflectances = []
    transmittances = []
    losses = []

    for i in range(nfreq):
        wl = 1 / flux_freqs[i]
        R = -bend_refl_flux[i] / straight_tran_flux[i]
        T = bend_tran_flux[i] / straight_tran_flux[i]
        L = 1 - R - T

        wavelengths.append(wl)
        reflectances.append(R)
        transmittances.append(T)
        losses.append(L)

    run.log(step=2, phase="bent_waveguide",
            max_transmittance=max(transmittances),
            min_transmittance=min(transmittances),
            max_reflectance=max(reflectances),
            max_loss=max(losses))

    # Create transmission plot
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(wavelengths, reflectances, "b-", linewidth=2, label="Reflectance")
    ax.plot(wavelengths, transmittances, "r-", linewidth=2, label="Transmittance")
    ax.plot(wavelengths, losses, "g-", linewidth=2, label="Loss")
    ax.set_xlim(5.0, 10.0)
    ax.set_ylim(0, 1)
    ax.set_xlabel("Wavelength (μm)")
    ax.set_ylabel("Coefficient")
    ax.set_title("Waveguide Bend Transmission Analysis")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()

    run.log_matplotlib("transmission_analysis", fig)
    plt.close(fig)

    run.log(step=3, peak_transmittance=max(transmittances),
            simulation_completed=True)

    print(f"📊 Peak transmittance: {max(transmittances):.3f}")
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

tran_fr = mp.FluxRegion(
    center=mp.Vector3(x, y, 0),
    size=mp.Vector3(0, 2*w, 0)  # Perpendicular to propagation
)
tran = sim.add_flux(fcen, df, nfreq, tran_fr)

# After straight simulation
straight_refl_data = sim.get_flux_data(refl)

# In bent simulation
sim.load_minus_flux_data(refl, straight_refl_data)

R = -bend_refl_flux[i] / straight_tran_flux[i]  # Negative because subtracted
T = bend_tran_flux[i] / straight_tran_flux[i]
L = 1 - R - T  # Conservation of energy