"""
Holey Waveguide Cavity with OptixLog Integration

Simulates transmission through a cavity formed by holes in a waveguide,
and finds resonant modes using Harminv analysis.

Based on the Meep tutorial: holey-wvg-cavity.py
Reference: S. Fan et al., J. Opt. Soc. Am. B, 12 (7), 1267-1272 (1995)
"""

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

api_key = os.getenv("OPTIX_API_KEY", "")
api_url = os.getenv("OPTIX_API_URL", "https://optixlog.com")
project_name = os.getenv("OPTIX_PROJECT", "MeepExamples")


def main(args):
    """Main simulation function for holey waveguide cavity."""
    
    if not optixlog.is_master_process():
        return
    
    try:
        run_type = "resonant_modes" if args.resonant_modes else "transmission"
        
        client = optixlog.init(
            api_key=api_key,
            api_url=api_url,
            project=project_name,
            run_name=f"holey_waveguide_cavity_{run_type}",
            config={
                "simulation_type": "cavity",
                "analysis_mode": run_type,
                "num_holes": args.N,
            },
            create_project_if_not_exists=True
        )
        
        # Simulation parameters
        resolution = 20
        eps = 13         # dielectric constant of waveguide
        w = 1.2          # width of waveguide
        r = 0.36         # radius of holes
        d = 1.4          # defect spacing
        N = args.N       # number of holes on either side
        
        sy = args.sy     # size of cell in y direction
        pad = 2          # padding between last hole and PML
        dpml = 1         # PML thickness
        
        sx = 2 * (pad + dpml + N) + d - 1
        
        fcen = args.fcen
        df = args.df
        
        # Log parameters
        client.log(
            step=0,
            resolution=resolution,
            waveguide_epsilon=eps,
            waveguide_width=w,
            hole_radius=r,
            defect_spacing=d,
            num_holes_per_side=N,
            center_frequency=fcen,
            frequency_width=df,
        )
        
        # Set up cell and geometry
        cell = mp.Vector3(sx, sy, 0)
        blk = mp.Block(size=mp.Vector3(mp.inf, w, mp.inf), material=mp.Medium(epsilon=eps))
        geometry = [blk]
        
        # Add holes
        for i in range(N):
            geometry.append(mp.Cylinder(r, center=mp.Vector3(d / 2 + i)))
            geometry.append(mp.Cylinder(r, center=mp.Vector3(-(d / 2 + i))))
        
        sim = mp.Simulation(
            cell_size=cell,
            geometry=geometry,
            sources=[],
            boundary_layers=[mp.PML(dpml)],
            resolution=resolution,
        )
        
        if args.resonant_modes:
            # Mode finding simulation
            sim.sources.append(
                mp.Source(
                    mp.GaussianSource(fcen, fwidth=df),
                    component=mp.Hz,
                    center=mp.Vector3()
                )
            )
            
            sim.symmetries.append(mp.Mirror(mp.Y, phase=-1))
            sim.symmetries.append(mp.Mirror(mp.X, phase=-1))
            
            harminv = mp.Harminv(mp.Hz, mp.Vector3(), fcen, df)
            
            sim.run(
                mp.at_beginning(mp.output_epsilon),
                mp.after_sources(harminv),
                until_after_sources=400,
            )
            
            modes = harminv.modes
            client.log(step=1, num_modes_found=len(modes))
            
            for i, mode in enumerate(modes):
                client.log(
                    step=2 + i,
                    mode_number=i + 1,
                    frequency=float(mode.freq),
                    quality_factor=float(mode.Q) if mode.Q > 0 else 0,
                )
                
        else:
            # Transmission spectrum simulation
            sim.sources.append(
                mp.Source(
                    mp.GaussianSource(fcen, fwidth=df),
                    component=mp.Ey,
                    center=mp.Vector3(-0.5 * sx + dpml),
                    size=mp.Vector3(0, w),
                )
            )
            
            sim.symmetries.append(mp.Mirror(mp.Y, phase=-1))
            
            freg = mp.FluxRegion(
                center=mp.Vector3(0.5 * sx - dpml - 0.5),
                size=mp.Vector3(0, 2 * w)
            )
            nfreq = 500
            trans = sim.add_flux(fcen, df, nfreq, freg)
            
            sim.run(
                mp.at_beginning(mp.output_epsilon),
                until_after_sources=mp.stop_when_fields_decayed(
                    50, mp.Ey, mp.Vector3(0.5 * sx - dpml - 0.5), 1e-3
                ),
            )
            
            flux_freqs = mp.get_flux_freqs(trans)
            flux_data = mp.get_fluxes(trans)
            
            client.log(
                step=2,
                num_frequencies=len(flux_freqs),
                max_transmission=float(np.max(flux_data)),
            )
        
    except Exception as e:
        print(f"Simulation Error: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-r", "--resonant_modes", action="store_true", default=False)
    parser.add_argument("-N", type=int, default=3)
    parser.add_argument("-sy", type=int, default=6)
    parser.add_argument("-fcen", type=float, default=0.25)
    parser.add_argument("-df", type=float, default=0.2)
    args = parser.parse_args()
    main(args)

# Set your OptixLog API key
export OPTIX_API_KEY="your_api_key_here"

# Find resonant modes with Harminv
python holey-wvg-cavity.py -r

# Compute transmission spectrum (default)
python holey-wvg-cavity.py

# Customize number of holes
python holey-wvg-cavity.py -N 5 -r

# Custom frequency range
python holey-wvg-cavity.py -fcen 0.3 -df 0.15