"""
Binary Grating Diffraction Analysis with OptixLog Integration

Computes the transmittance of diffraction orders for a binary grating.
"""

import os
import math
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)

    # Grating parameters
    resolution = 60
    dpml = 1.0
    dsub = 3.0
    dpad = 3.0
    gp = 10.0  # grating period
    gh = 0.5   # grating height
    gdc = 0.5  # duty cycle
    glass_index = 1.5

    # Wavelength range
    wvl_min = 0.4
    wvl_max = 0.6
    fmin = 1 / wvl_max
    fmax = 1 / wvl_min
    fcen = 0.5 * (fmin + fmax)
    df = fmax - fmin

    # Cell size
    sx = dpml + dsub + gh + dpad + dpml
    sy = gp
    cell_size = mp.Vector3(sx, sy, 0)
    
    nfreq = 21
    nmode = 10  # number of diffraction orders to compute

    run = project.run(
        name="binary_grating_diffraction",
        config={
            "simulation_type": "binary_grating",
            "resolution": resolution,
            "grating_period": gp,
            "grating_height": gh,
            "grating_duty_cycle": gdc,
            "glass_index": glass_index,
            "wavelength_min": wvl_min,
            "wavelength_max": wvl_max,
            "num_modes": nmode
        }
    )
    
    run.log(step=0, grating_period=gp, grating_height=gh)

    pml_layers = [mp.PML(thickness=dpml, direction=mp.X)]
    glass = mp.Medium(index=glass_index)
    symmetries = [mp.Mirror(mp.Y)]

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

    # ===== Empty cell simulation =====
    print("⚡ Running empty cell simulation...")
    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell_size,
        boundary_layers=pml_layers,
        k_point=mp.Vector3(0, 0, 0),
        default_material=glass,
        sources=sources,
        symmetries=symmetries,
    )

    mon_pt = mp.Vector3(0.5 * sx - dpml - 0.5 * dpad, 0, 0)
    flux_mon = sim.add_flux(
        fcen, df, nfreq, mp.FluxRegion(center=mon_pt, size=mp.Vector3(0, sy, 0))
    )

    sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, mon_pt, 1e-9))
    input_flux = mp.get_fluxes(flux_mon)

    run.log(step=1, phase="empty_cell", max_input_flux=float(max(input_flux)))

    sim.reset_meep()

    # ===== Grating simulation =====
    print("⚡ Running grating simulation...")
    geometry = [
        mp.Block(
            material=glass,
            size=mp.Vector3(dpml + dsub, mp.inf, mp.inf),
            center=mp.Vector3(-0.5 * sx + 0.5 * (dpml + dsub), 0, 0),
        ),
        mp.Block(
            material=glass,
            size=mp.Vector3(gh, gdc * gp, mp.inf),
            center=mp.Vector3(-0.5 * sx + dpml + dsub + 0.5 * gh, 0, 0),
        ),
    ]

    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell_size,
        boundary_layers=pml_layers,
        geometry=geometry,
        k_point=mp.Vector3(0, 0, 0),
        sources=sources,
        symmetries=symmetries,
    )

    mode_mon = sim.add_flux(
        fcen, df, nfreq, mp.FluxRegion(center=mon_pt, size=mp.Vector3(0, sy, 0))
    )

    sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, mon_pt, 1e-9))

    # Get eigenmode coefficients
    freqs = mp.get_eigenmode_freqs(mode_mon)
    res = sim.get_eigenmode_coefficients(
        mode_mon, range(1, nmode + 1), eig_parity=mp.ODD_Z + mp.EVEN_Y
    )
    coeffs = res.alpha
    kdom = res.kdom

    # Process diffraction data
    print("📊 Processing diffraction orders...")
    mode_wvl = []
    mode_angle = []
    mode_tran = []

    for nm in range(nmode):
        for nf in range(nfreq):
            wvl = 1 / freqs[nf]
            angle = math.degrees(math.acos(kdom[nm * nfreq + nf].x / freqs[nf]))
            tran = abs(coeffs[nm, nf, 0]) ** 2 / input_flux[nf]
            tran = 0.5 * tran if nm != 0 else tran  # Symmetry factor
            
            mode_wvl.append(wvl)
            mode_angle.append(angle)
            mode_tran.append(tran)

    run.log(step=2, phase="grating_analysis",
            num_diffraction_orders=nmode,
            max_transmission=float(max(mode_tran)))

    # Create diffraction map
    fig, ax = plt.subplots(figsize=(12, 8))
    tran_max = round(max(mode_tran), 1)
    
    im = ax.pcolormesh(
        np.reshape(mode_wvl, (nmode, nfreq)),
        np.reshape(mode_angle, (nmode, nfreq)),
        np.reshape(mode_tran, (nmode, nfreq)),
        cmap="Blues",
        shading="nearest",
        vmin=0,
        vmax=tran_max,
    )
    ax.set_xlim(min(mode_wvl), max(mode_wvl))
    ax.set_ylim(min(mode_angle), max(mode_angle))
    ax.set_xlabel("Wavelength (μm)")
    ax.set_ylabel("Diffraction Angle (degrees)")
    ax.set_title("Binary Grating: Transmittance of Diffraction Orders")
    plt.colorbar(im, ax=ax, label="Transmittance")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    
    run.log_matplotlib("diffraction_orders", fig)
    plt.close(fig)

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

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


if __name__ == "__main__":
    main()

sim = mp.Simulation(
    ...
    k_point=mp.Vector3(0, 0, 0),  # Bloch periodic
    ...
)

res = sim.get_eigenmode_coefficients(
    mode_mon, 
    range(1, nmode + 1),  # Mode indices
    eig_parity=mp.ODD_Z + mp.EVEN_Y
)
coeffs = res.alpha  # Complex amplitudes
kdom = res.kdom     # Dominant k-vector for each mode

angle = math.degrees(math.acos(kdom[idx].x / freq))