"""
Holey Waveguide Band Structure with OptixLog Integration

Computes the band structure for a photonic crystal waveguide
formed by a periodic sequence of holes in a dielectric 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
    eps = 13   # waveguide permittivity
    w = 1.2    # waveguide width
    r = 0.36   # hole radius
    sy = 12    # cell size in y
    dpml = 1
    fcen = 0.25
    df = 1.5
    resolution = 20
    
    run = project.run(
        name="holey_waveguide_bands",
        config={
            "simulation_type": "band_structure",
            "waveguide_epsilon": eps,
            "waveguide_width": w,
            "hole_radius": r,
            "resolution": resolution
        }
    )
    
    run.log(step=0, waveguide_epsilon=eps, waveguide_width=w, hole_radius=r)
    
    # Cell (1 period in x, extended in y)
    cell = mp.Vector3(1, sy)
    
    # Geometry: waveguide + hole
    geometry = [
        mp.Block(size=mp.Vector3(mp.inf, w, mp.inf), material=mp.Medium(epsilon=eps)),
        mp.Cylinder(radius=r)
    ]
    
    # Broadband source for mode excitation
    sources = [
        mp.Source(
            src=mp.GaussianSource(fcen, fwidth=df),
            component=mp.Hz,
            center=mp.Vector3(0.1234),  # Slightly off-center
        )
    ]
    
    # Symmetry (odd Hz modes)
    symmetries = [mp.Mirror(direction=mp.Y, phase=-1)]
    
    sim = mp.Simulation(
        cell_size=cell,
        geometry=geometry,
        sources=sources,
        symmetries=symmetries,
        boundary_layers=[mp.PML(dpml, direction=mp.Y)],
        resolution=resolution,
    )
    
    # K-point sweep
    k_interp = 19
    k_points = mp.interpolate(k_interp, [mp.Vector3(), mp.Vector3(0.5)])
    
    print("⚡ Running k-point sweep...")
    all_freqs = sim.run_k_points(300, k_points)
    
    run.log(step=1, simulation_completed=True, num_k_points=len(k_points))
    
    # Process results
    k_values = []
    freq_data = []
    
    for i, (freqs, kpt) in enumerate(zip(all_freqs, k_points)):
        k_values.append(kpt.x)
        freq_list = [f.real for f in freqs]
        freq_data.append(freq_list)
        
        if len(freq_list) > 0:
            run.log(step=2 + i,
                    k_point=float(kpt.x),
                    num_modes=len(freq_list),
                    min_frequency=float(min(freq_list)) if freq_list else 0,
                    max_frequency=float(max(freq_list)) if freq_list else 0)
    
    # Create band structure plot
    fig, ax = plt.subplots(figsize=(10, 8))
    
    for k, freqs in zip(k_values, freq_data):
        if freqs:
            ax.scatter([k] * len(freqs), freqs, c='blue', s=15)
    
    # Add light line
    light_line_k = np.linspace(0, 0.5, 100)
    light_line_f = light_line_k  # ω = ck in normalized units
    ax.plot(light_line_k, light_line_f, 'r--', linewidth=1.5, label='Light line')
    
    ax.set_xlabel('Wave vector k (2π/a)', fontsize=12)
    ax.set_ylabel('Frequency (c/a)', fontsize=12)
    ax.set_title('Holey Waveguide Band Structure', fontsize=14)
    ax.set_xlim(0, 0.5)
    ax.set_ylim(0, 0.8)
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    run.log_matplotlib("band_structure", fig)
    plt.close(fig)
    
    # Summary statistics
    all_freqs_flat = [f for freqs in freq_data for f in freqs if f > 0]
    if all_freqs_flat:
        run.log(step=100,
                total_modes_found=len(all_freqs_flat),
                frequency_range_min=float(min(all_freqs_flat)),
                frequency_range_max=float(max(all_freqs_flat)),
                analysis_complete=True)
    
    print(f"📊 Found {len(all_freqs_flat)} modes")
    print(f"\n✅ Band structure complete!")


if __name__ == "__main__":
    main()

k_points = mp.interpolate(k_interp, [mp.Vector3(), mp.Vector3(0.5)])
all_freqs = sim.run_k_points(300, k_points)

# Light line: ω = ck
light_line_f = light_line_k  # In normalized units