"""
Ring Resonator Simulation with OptixLog Integration

Simulates an optical ring resonator coupled to a straight waveguide.
"""

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)
    
    # Parameters
    resolution = 10
    n = 3.4  # waveguide index (silicon)
    w = 1    # waveguide width
    r = 1    # ring radius
    pad = 4  # padding
    dpml = 2
    
    # Ring parameters
    sxy = 2 * (r + w + pad + dpml)
    
    # Source parameters
    fcen = 0.15
    df = 0.1
    nfreq = 500
    
    run = project.run(
        name="ring_resonator_simulation",
        config={
            "simulation_type": "ring_resonator",
            "resolution": resolution,
            "ring_radius": r,
            "waveguide_width": w,
            "refractive_index": n,
            "center_frequency": fcen,
            "num_frequencies": nfreq
        }
    )
    
    run.log(step=0, ring_radius=r, waveguide_width=w, refractive_index=n)
    
    cell = mp.Vector3(sxy, sxy)
    pml_layers = [mp.PML(dpml)]
    
    # Geometry: ring + bus waveguide
    geometry = [
        # Bus waveguide (straight)
        mp.Block(
            center=mp.Vector3(0, r + w/2 + w/2),
            size=mp.Vector3(mp.inf, w, mp.inf),
            material=mp.Medium(index=n)
        ),
        # Ring (cylinder with hole)
        mp.Cylinder(
            radius=r + w,
            height=mp.inf,
            material=mp.Medium(index=n)
        ),
        mp.Cylinder(
            radius=r,
            height=mp.inf,
            material=mp.Medium(index=1)  # Air hole
        ),
    ]
    
    # Source in bus waveguide
    sources = [
        mp.Source(
            mp.GaussianSource(fcen, fwidth=df),
            component=mp.Ez,
            center=mp.Vector3(-r - w - pad/2, r + w/2 + w/2),
            size=mp.Vector3(0, w),
        )
    ]
    
    sim = mp.Simulation(
        cell_size=cell,
        geometry=geometry,
        sources=sources,
        boundary_layers=pml_layers,
        resolution=resolution,
    )
    
    # Transmission monitor (through port)
    tran_fr = mp.FluxRegion(
        center=mp.Vector3(r + w + pad/2, r + w/2 + w/2),
        size=mp.Vector3(0, 2*w)
    )
    tran = sim.add_flux(fcen, df, nfreq, tran_fr)
    
    # Run straight waveguide for normalization
    print("⚡ Running straight waveguide reference...")
    sim_straight = mp.Simulation(
        cell_size=cell,
        geometry=[
            mp.Block(
                center=mp.Vector3(0, r + w/2 + w/2),
                size=mp.Vector3(mp.inf, w, mp.inf),
                material=mp.Medium(index=n)
            )
        ],
        sources=sources,
        boundary_layers=pml_layers,
        resolution=resolution,
    )
    
    tran_straight = sim_straight.add_flux(fcen, df, nfreq, tran_fr)
    
    pt = mp.Vector3(r + w + pad/2, r + w/2 + w/2)
    sim_straight.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, pt, 1e-3))
    
    straight_flux = mp.get_fluxes(tran_straight)
    flux_freqs = mp.get_flux_freqs(tran_straight)
    
    run.log(step=1, phase="straight_reference_completed")
    
    # Run ring resonator
    print("⚡ Running ring resonator simulation...")
    sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, pt, 1e-3))
    
    ring_flux = mp.get_fluxes(tran)
    
    # Calculate transmission
    wavelengths = [1/f for f in flux_freqs]
    transmission = [ring_flux[i] / straight_flux[i] if straight_flux[i] > 0 else 0 
                   for i in range(nfreq)]
    
    run.log(step=2, phase="ring_simulation_completed",
            min_transmission=float(min(transmission)),
            max_transmission=float(max(transmission)))
    
    # Find resonances (transmission dips)
    resonance_indices = []
    for i in range(1, len(transmission) - 1):
        if transmission[i] < transmission[i-1] and transmission[i] < transmission[i+1]:
            if transmission[i] < 0.8:  # Significant dip
                resonance_indices.append(i)
    
    resonance_wavelengths = [wavelengths[i] for i in resonance_indices]
    resonance_depths = [1 - transmission[i] for i in resonance_indices]
    
    # Calculate FSR if multiple resonances found
    if len(resonance_wavelengths) >= 2:
        fsr = abs(resonance_wavelengths[0] - resonance_wavelengths[1])
    else:
        fsr = 0
    
    run.log(step=3,
            num_resonances=len(resonance_indices),
            free_spectral_range=fsr)
    
    # Create plots
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # Transmission spectrum
    axes[0, 0].plot(wavelengths, transmission, 'b-', linewidth=1.5)
    for idx in resonance_indices:
        axes[0, 0].axvline(wavelengths[idx], color='r', linestyle='--', alpha=0.5)
    axes[0, 0].set_xlabel('Wavelength (μm)')
    axes[0, 0].set_ylabel('Transmission')
    axes[0, 0].set_title('Ring Resonator Transmission Spectrum')
    axes[0, 0].set_ylim(0, 1.1)
    axes[0, 0].grid(True, alpha=0.3)
    
    # Frequency spectrum
    axes[0, 1].plot(flux_freqs, transmission, 'g-', linewidth=1.5)
    axes[0, 1].set_xlabel('Frequency (c/a)')
    axes[0, 1].set_ylabel('Transmission')
    axes[0, 1].set_title('Transmission vs Frequency')
    axes[0, 1].grid(True, alpha=0.3)
    
    # Resonance depths
    if resonance_wavelengths:
        axes[1, 0].bar(range(len(resonance_depths)), resonance_depths, color='red', alpha=0.7)
        axes[1, 0].set_xlabel('Resonance Index')
        axes[1, 0].set_ylabel('Extinction (1-T)')
        axes[1, 0].set_title('Resonance Extinction Depths')
        axes[1, 0].grid(True, alpha=0.3)
    
    # Loss (1-T) on log scale
    loss = [1 - t if t < 1 else 1e-6 for t in transmission]
    axes[1, 1].semilogy(wavelengths, loss, 'purple', linewidth=1.5)
    axes[1, 1].set_xlabel('Wavelength (μm)')
    axes[1, 1].set_ylabel('1 - Transmission')
    axes[1, 1].set_title('Loss Spectrum (Log Scale)')
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    run.log_matplotlib("ring_resonator_analysis", fig)
    plt.close(fig)
    
    run.log(step=4, simulation_completed=True)
    
    print(f"📊 Found {len(resonance_indices)} resonances")
    if fsr > 0:
        print(f"📊 FSR: {fsr:.4f} μm")
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

geometry = [
    mp.Cylinder(radius=r + w, material=mp.Medium(index=n)),
    mp.Cylinder(radius=r, material=mp.Medium(index=1)),  # Air core
]

for i in range(1, len(transmission) - 1):
    if transmission[i] < transmission[i-1] and transmission[i] < transmission[i+1]:
        if transmission[i] < 0.8:
            resonance_indices.append(i)

transmission = ring_flux / straight_flux