"""
Mode Decomposition Analysis with OptixLog Integration

Analyzes waveguide taper reflectance using eigenmode decomposition.
"""

import os
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", "")
project_name = os.getenv("OPTIX_PROJECT", "MeepExamples")


def main():
    if not optixlog.is_master_process():
        return
    
    try:
        client = optixlog.init(
            api_key=api_key,
            project=project_name,
            run_name="mode_decomposition_analysis",
            config={"simulation_type": "mode_decomposition"},
            create_project_if_not_exists=True
        )
        
        resolution = 25
        w1, w2 = 1.0, 2.0
        Lw = 10.0
        dair, dpml_x, dpml_y = 3.0, 6.0, 2.0
        sy = dpml_y + dair + w2 + dair + dpml_y
        
        Si = mp.Medium(epsilon=12.0)
        lcen, fcen = 6.67, 1/6.67
        
        Lts = [2**m for m in range(4)]  # 1, 2, 4, 8
        
        client.log(step=0, resolution=resolution, input_width=w1, output_width=w2, num_taper_lengths=len(Lts))
        
        R_coeffs, R_flux = [], []
        
        for Lt in Lts:
            sx = dpml_x + Lw + Lt + Lw + dpml_x
            cell_size = mp.Vector3(sx, sy, 0)
            boundary_layers = [mp.PML(dpml_x, direction=mp.X), mp.PML(dpml_y, direction=mp.Y)]
            
            src_pt = mp.Vector3(-0.5 * sx + dpml_x + 0.2 * Lw)
            sources = [mp.EigenModeSource(src=mp.GaussianSource(fcen, fwidth=0.2 * fcen), center=src_pt, size=mp.Vector3(y=sy - 2 * dpml_y), eig_parity=mp.ODD_Z + mp.EVEN_Y)]
            
            # Straight waveguide reference
            vertices = [mp.Vector3(-0.5 * sx - 1, 0.5 * w1), mp.Vector3(0.5 * sx + 1, 0.5 * w1), mp.Vector3(0.5 * sx + 1, -0.5 * w1), mp.Vector3(-0.5 * sx - 1, -0.5 * w1)]
            
            sim = mp.Simulation(resolution=resolution, cell_size=cell_size, boundary_layers=boundary_layers, geometry=[mp.Prism(vertices, height=mp.inf, material=Si)], sources=sources, symmetries=[mp.Mirror(mp.Y)])
            
            mon_pt = mp.Vector3(-0.5 * sx + dpml_x + 0.7 * Lw)
            flux = sim.add_flux(fcen, 0, 1, mp.FluxRegion(center=mon_pt, size=mp.Vector3(y=sy - 2 * dpml_y)))
            
            sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, mon_pt, 1e-9))
            
            res = sim.get_eigenmode_coefficients(flux, [1], eig_parity=mp.ODD_Z + mp.EVEN_Y)
            incident_coeffs = res.alpha
            incident_flux = mp.get_fluxes(flux)
            incident_flux_data = sim.get_flux_data(flux)
            
            sim.reset_meep()
            
            # Tapered waveguide
            vertices = [mp.Vector3(-0.5 * sx - 1, 0.5 * w1), mp.Vector3(-0.5 * Lt, 0.5 * w1), mp.Vector3(0.5 * Lt, 0.5 * w2), mp.Vector3(0.5 * sx + 1, 0.5 * w2), mp.Vector3(0.5 * sx + 1, -0.5 * w2), mp.Vector3(0.5 * Lt, -0.5 * w2), mp.Vector3(-0.5 * Lt, -0.5 * w1), mp.Vector3(-0.5 * sx - 1, -0.5 * w1)]
            
            sim = mp.Simulation(resolution=resolution, cell_size=cell_size, boundary_layers=boundary_layers, geometry=[mp.Prism(vertices, height=mp.inf, material=Si)], sources=sources, symmetries=[mp.Mirror(mp.Y)])
            
            flux = sim.add_flux(fcen, 0, 1, mp.FluxRegion(center=mon_pt, size=mp.Vector3(y=sy - 2 * dpml_y)))
            sim.load_minus_flux_data(flux, incident_flux_data)
            
            sim.run(until_after_sources=mp.stop_when_fields_decayed(50, mp.Ez, mon_pt, 1e-9))
            
            res2 = sim.get_eigenmode_coefficients(flux, [1], eig_parity=mp.ODD_Z + mp.EVEN_Y)
            taper_coeffs = res2.alpha
            taper_flux = mp.get_fluxes(flux)
            
            R_coeffs.append(abs(taper_coeffs[0, 0, 1]) ** 2 / abs(incident_coeffs[0, 0, 0]) ** 2)
            R_flux.append(-taper_flux[0] / incident_flux[0])
            
            client.log(step=len(R_coeffs), taper_length=Lt, R_mode_decomposition=float(R_coeffs[-1]), R_poynting_flux=float(R_flux[-1]))
        
        client.log(step=100, simulation_completed=True)
        
    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()