"""
Binary Grating Phase Map Simulation with OptixLog Integration

This script simulates binary grating transmission and phase characteristics using Meep
and logs the results to OptixLog for visualization and tracking.

Usage:
    # Set your OptixLog API key
    export OPTIX_API_KEY="proj_your_api_key_here"
    
    # Run the simulation
    python binary_grating_phasemap.py                    # Default parameters
    python binary_grating_phasemap.py -gp 0.4 -gh 0.7    # Custom grating parameters
    python binary_grating_phasemap.py -oddz              # Odd Z symmetry
"""

import argparse
import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'Coupler', 'sdk'))

import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import numpy as np
import numpy.matlib

import meep as mp
import optixlog

resolution = 60  # pixels/μm

dpml = 1.0  # PML thickness
dsub = 3.0  # substrate thickness
dpad = 3.0  # padding between grating and PML

wvl_min = 0.4  # min wavelength
wvl_max = 0.6  # max wavelength
fmin = 1 / wvl_max  # min frequency
fmax = 1 / wvl_min  # max frequency
fcen = 0.5 * (fmin + fmax)  # center frequency
df = fmax - fmin  # frequency width
nfreq = 21  # number of frequency bins

k_point = mp.Vector3(0, 0, 0)

glass = mp.Medium(index=1.5)


def grating(gp, gh, gdc, oddz):
    sx = dpml + dsub + gh + dpad + dpml
    sy = gp

    cell_size = mp.Vector3(sx, sy, 0)
    pml_layers = [mp.PML(thickness=dpml, direction=mp.X)]

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

    symmetries = [mp.Mirror(mp.Y, phase=+1 if oddz else -1)]

    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell_size,
        boundary_layers=pml_layers,
        k_point=k_point,
        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=100)

    input_flux = mp.get_fluxes(flux_mon)

    sim.reset_meep()

    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=k_point,
        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=300)

    freqs = mp.get_eigenmode_freqs(mode_mon)
    res = sim.get_eigenmode_coefficients(
        mode_mon, [1], eig_parity=mp.ODD_Z + mp.EVEN_Y if oddz else mp.EVEN_Z + mp.ODD_Y
    )
    coeffs = res.alpha

    mode_wvl = [1 / freqs[nf] for nf in range(nfreq)]
    mode_tran = [abs(coeffs[0, nf, 0]) ** 2 / input_flux[nf] for nf in range(nfreq)]
    mode_phase = [np.angle(coeffs[0, nf, 0]) for nf in range(nfreq)]

    return mode_wvl, mode_tran, mode_phase


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Binary Grating Phase Map Simulation with OptixLog Integration"
    )
    parser.add_argument("-gp", type=float, default=0.35, help="grating periodicity (μm)")
    parser.add_argument("-gh", type=float, default=0.6, help="grating height (μm)")
    parser.add_argument("-oddz", action="store_true", default=False, help="use odd Z symmetry")
    args = parser.parse_args()

    # Initialize OptixLog client
    config = {
        "simulation_type": "binary_grating_phasemap",
        "grating_period": args.gp,
        "grating_height": args.gh,
        "odd_z_symmetry": args.oddz,
        "wavelength_range": [wvl_min, wvl_max],
    }

    client = optixlog.init(
        api_key=os.getenv("OPTIX_API_KEY"),
        project="Binary Grating Phase Map",
        run_name=f"gp{args.gp}_gh{args.gh}",
        config=config,
        create_project_if_not_exists=True
    )

    # Define duty cycle range
    gdc = np.arange(0.1, 1.0, 0.1)

    mode_tran = np.empty((gdc.size, nfreq))
    mode_phase = np.empty((gdc.size, nfreq))
    
    for n in range(gdc.size):
        mode_wvl, mode_tran[n, :], mode_phase[n, :] = grating(
            args.gp, args.gh, gdc[n], args.oddz
        )
        
        client.log(n + 1,
            duty_cycle=gdc[n],
            max_transmittance=float(np.max(mode_tran[n, :])),
            phase_range=float(np.max(mode_phase[n, :]) - np.min(mode_phase[n, :]))
        )

    # Create phase map visualization
    plt.figure(figsize=(16, 8), dpi=150)

    plt.subplot(1, 2, 1)
    plt.pcolormesh(mode_wvl, gdc, mode_tran, cmap="hot_r", shading="gouraud")
    plt.xlabel("wavelength (μm)")
    plt.ylabel("grating duty cycle")
    plt.title("Transmittance")
    plt.colorbar()

    plt.subplot(1, 2, 2)
    plt.pcolormesh(mode_wvl, gdc, mode_phase, cmap="RdBu", shading="gouraud")
    plt.xlabel("wavelength (μm)")
    plt.ylabel("grating duty cycle")
    plt.title("Phase (radians)")
    plt.colorbar()

    plt.tight_layout()
    plt.savefig("binary_grating_phasemap.png", dpi=150, bbox_inches="tight")

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

# Run with default parameters
python binary_grating_phasemap.py

# Custom grating parameters
python binary_grating_phasemap.py -gp 0.4 -gh 0.7

# Use odd Z symmetry (TM polarization)
python binary_grating_phasemap.py -oddz