"""
GDSII Coupler Simulation with OptixLog Integration

Computes the S-parameters of a directional coupler.
"""

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


# Material definitions
SILICON = mp.Medium(epsilon=12)
OXIDE = mp.Medium(epsilon=2.25)

# Parameters
T_SI = 0.22   # Silicon thickness
T_OXIDE = 1.0 # Oxide thickness
DPML = 1      # PML thickness
FCEN = 1 / 1.55  # Center frequency (1.55 μm)
DF = 0.2 * FCEN


def create_coupler_geometry(d=0.1, si_zmin=-0.11, si_zmax=0.11):
    """Create directional coupler geometry"""
    # Simple 2D coupler for demonstration
    # In practice, load from GDSII
    w = 0.5  # waveguide width
    L = 10   # coupler length
    gap = d  # separation
    
    geometry = [
        # Upper branch
        mp.Block(
            center=mp.Vector3(0, w/2 + gap/2),
            size=mp.Vector3(L, w, mp.inf),
            material=SILICON
        ),
        # Lower branch
        mp.Block(
            center=mp.Vector3(0, -w/2 - gap/2),
            size=mp.Vector3(L, w, mp.inf),
            material=SILICON
        ),
    ]
    
    return geometry, L, w


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 = 50
    d = 0.1  # branch separation
    
    run = project.run(
        name=f"directional_coupler_d{d:.2f}",
        config={
            "simulation_type": "directional_coupler",
            "resolution": resolution,
            "branch_separation": d,
            "wavelength": 1.55,
            "silicon_thickness": T_SI
        }
    )
    
    run.log(step=0, resolution=resolution, branch_separation=d)
    
    # Create geometry
    geometry, L, w = create_coupler_geometry(d)
    
    # Cell size
    sx = L + 2 * DPML + 4
    sy = 4 * w + 2 * d + 2 * DPML
    cell = mp.Vector3(sx, sy)
    
    # Source at input port
    sources = [
        mp.EigenModeSource(
            src=mp.GaussianSource(FCEN, fwidth=DF),
            center=mp.Vector3(-L/2 - 1, w/2 + d/2),
            size=mp.Vector3(0, 2*w),
            eig_parity=mp.EVEN_Y + mp.ODD_Z,
        )
    ]
    
    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell,
        boundary_layers=[mp.PML(DPML)],
        sources=sources,
        geometry=geometry,
    )
    
    # Mode monitors at each port
    port_through = sim.add_mode_monitor(
        FCEN, 0, 1,
        mp.ModeRegion(center=mp.Vector3(L/2 + 1, w/2 + d/2), size=mp.Vector3(0, 2*w))
    )
    port_cross = sim.add_mode_monitor(
        FCEN, 0, 1,
        mp.ModeRegion(center=mp.Vector3(L/2 + 1, -w/2 - d/2), size=mp.Vector3(0, 2*w))
    )
    port_input = sim.add_mode_monitor(
        FCEN, 0, 1,
        mp.ModeRegion(center=mp.Vector3(-L/2 - 1, w/2 + d/2), size=mp.Vector3(0, 2*w))
    )
    
    print("⚡ Running simulation...")
    sim.run(until_after_sources=100)
    
    run.log(step=1, simulation_phase="completed")
    
    # Calculate S-parameters
    eig_parity = mp.EVEN_Y + mp.ODD_Z
    
    input_coeff = sim.get_eigenmode_coefficients(port_input, [1], eig_parity=eig_parity).alpha[0, 0, 0]
    through_coeff = sim.get_eigenmode_coefficients(port_through, [1], eig_parity=eig_parity).alpha[0, 0, 0]
    cross_coeff = sim.get_eigenmode_coefficients(port_cross, [1], eig_parity=eig_parity).alpha[0, 0, 0]
    
    # Transmittance (power ratios)
    S21 = abs(through_coeff) ** 2 / abs(input_coeff) ** 2
    S31 = abs(cross_coeff) ** 2 / abs(input_coeff) ** 2
    total = S21 + S31
    
    # Coupling ratio
    coupling_ratio = S31 / (S21 + S31) if (S21 + S31) > 0 else 0
    
    print(f"📊 S21 (through): {S21:.4f}")
    print(f"📊 S31 (cross): {S31:.4f}")
    print(f"📊 Total: {total:.4f}")
    print(f"📊 Coupling ratio: {coupling_ratio:.4f}")
    
    run.log(step=2,
            S21_transmittance=S21,
            S31_transmittance=S31,
            total_transmittance=total,
            coupling_ratio=coupling_ratio,
            simulation_completed=True)
    
    # Create visualization
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # Bar chart of S-parameters
    params = ['S21 (through)', 'S31 (cross)', 'Loss']
    values = [S21, S31, 1 - total]
    colors = ['blue', 'green', 'red']
    axes[0].bar(params, values, color=colors, alpha=0.7)
    axes[0].set_ylabel('Power Ratio')
    axes[0].set_title(f'Directional Coupler S-Parameters (gap={d} μm)')
    axes[0].set_ylim(0, 1)
    axes[0].grid(True, alpha=0.3)
    
    # Field pattern
    sim.plot2D(ax=axes[1], fields=mp.Ez)
    axes[1].set_title('Ez Field Distribution')
    
    plt.tight_layout()
    run.log_matplotlib("coupler_analysis", fig)
    plt.close(fig)
    
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

sources = [
    mp.EigenModeSource(
        src=mp.GaussianSource(FCEN, fwidth=DF),
        center=mp.Vector3(x, y),
        size=mp.Vector3(0, height),
        eig_parity=mp.EVEN_Y + mp.ODD_Z,
    )
]

port = sim.add_mode_monitor(
    fcen, 0, 1,
    mp.ModeRegion(center=center, size=size)
)

coeff = sim.get_eigenmode_coefficients(port, [1], eig_parity=parity)
S_param = abs(coeff.alpha[0, 0, 0]) ** 2 / abs(input) ** 2