"""
Zone Plate Simulation with OptixLog Integration

Computes the diffraction spectra of a binary-phase zone plate
in cylindrical coordinates and analyzes the focal properties.

Based on the Meep tutorial: zone_plate.py
"""

import os
import math
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", "")
api_url = os.getenv("OPTIX_API_URL", "https://optixlog.com")
project_name = os.getenv("OPTIX_PROJECT", "MeepExamples")


def main():
    """Main simulation function for zone plate analysis."""
    
    if not optixlog.is_master_process():
        return
    
    try:
        client = optixlog.init(
            api_key=api_key,
            api_url=api_url,
            project=project_name,
            run_name="zone_plate_simulation",
            config={
                "simulation_type": "zone_plate",
                "description": "Binary-phase zone plate focusing",
            },
            create_project_if_not_exists=True
        )
        
        # Simulation parameters
        resolution_um = 25
        pml_um = 1.0
        substrate_um = 2.0
        padding_um = 2.0
        height_um = 0.5
        focal_length_um = 200
        scan_length_z_um = 100
        farfield_resolution_um = 10
        
        wavelength_um = 0.5
        frequency = 1 / wavelength_um
        frequency_width = 0.2 * frequency
        
        num_zones = 25
        
        client.log(
            step=0,
            resolution=resolution_um,
            wavelength=wavelength_um,
            num_zones=num_zones,
            focal_length=focal_length_um,
        )
        
        pml_layers = [mp.PML(thickness=pml_um)]
        
        # Calculate zone radii
        zone_radius_um = np.zeros(num_zones)
        for n in range(1, num_zones + 1):
            zone_radius_um[n - 1] = math.sqrt(
                n * wavelength_um * (focal_length_um + n * wavelength_um / 4)
            )
        
        client.log(
            step=1,
            outermost_zone_radius=float(zone_radius_um[-1]),
            innermost_zone_radius=float(zone_radius_um[0]),
        )
        
        # Set up cell
        size_r_um = zone_radius_um[-1] + padding_um + pml_um
        size_z_um = pml_um + substrate_um + height_um + padding_um + pml_um
        cell_size = mp.Vector3(size_r_um, 0, size_z_um)
        
        # Circularly-polarized planewave source
        sources = [
            mp.Source(
                mp.GaussianSource(frequency, fwidth=frequency_width, is_integrated=True),
                component=mp.Er,
                center=mp.Vector3(0.5 * size_r_um, 0, -0.5 * size_z_um + pml_um),
                size=mp.Vector3(size_r_um),
            ),
            mp.Source(
                mp.GaussianSource(frequency, fwidth=frequency_width, is_integrated=True),
                component=mp.Ep,
                center=mp.Vector3(0.5 * size_r_um, 0, -0.5 * size_z_um + pml_um),
                size=mp.Vector3(size_r_um),
                amplitude=-1j,
            ),
        ]
        
        glass = mp.Medium(index=1.5)
        
        # Substrate
        geometry = [
            mp.Block(
                material=glass,
                size=mp.Vector3(size_r_um, 0, pml_um + substrate_um),
                center=mp.Vector3(0.5 * size_r_um, 0, -0.5 * size_z_um + 0.5 * (pml_um + substrate_um)),
            )
        ]
        
        # Zone plate rings
        zone_z = -0.5 * size_z_um + pml_um + substrate_um + 0.5 * height_um
        for n in range(num_zones - 1, -1, -1):
            geometry.append(
                mp.Block(
                    material=glass if n % 2 == 0 else mp.vacuum,
                    size=mp.Vector3(zone_radius_um[n], 0, height_um),
                    center=mp.Vector3(0.5 * zone_radius_um[n], 0, zone_z),
                )
            )
        
        sim = mp.Simulation(
            cell_size=cell_size,
            boundary_layers=pml_layers,
            resolution=resolution_um,
            sources=sources,
            geometry=geometry,
            dimensions=mp.CYLINDRICAL,
            m=-1,
        )
        
        # Near-to-far field monitor
        n2f_monitor = sim.add_near2far(
            frequency, 0, 1,
            mp.Near2FarRegion(
                center=mp.Vector3(0.5 * (size_r_um - pml_um), 0, 0.5 * size_z_um - pml_um),
                size=mp.Vector3(size_r_um - pml_um, 0, 0),
            ),
            mp.Near2FarRegion(
                center=mp.Vector3(size_r_um - pml_um, 0, 0.5 * size_z_um - pml_um - 0.5 * (height_um + padding_um)),
                size=mp.Vector3(0, 0, height_um + padding_um),
            ),
        )
        
        sim.run(
            until_after_sources=mp.stop_when_fields_decayed(
                50.0, mp.Er, mp.Vector3(0.5 * size_r_um, 0, 0), 1e-6
            )
        )
        
        # Compute far fields
        farfields_r = sim.get_farfields(
            n2f_monitor,
            farfield_resolution_um,
            center=mp.Vector3(
                0.5 * (size_r_um - pml_um), 0,
                -0.5 * size_z_um + pml_um + substrate_um + height_um + focal_length_um
            ),
            size=mp.Vector3(size_r_um - pml_um, 0, 0),
        )
        
        intensity_r = (
            np.absolute(farfields_r["Ex"])**2 +
            np.absolute(farfields_r["Ey"])**2 +
            np.absolute(farfields_r["Ez"])**2
        )
        
        # Find FWHM
        peak_idx = np.argmax(intensity_r)
        half_max = intensity_r[peak_idx] / 2
        fwhm_indices = np.where(intensity_r > half_max)[0]
        fwhm = (fwhm_indices[-1] - fwhm_indices[0]) / farfield_resolution_um if len(fwhm_indices) > 1 else 0
        
        client.log(
            step=4,
            max_intensity=float(np.max(intensity_r)),
            focal_spot_fwhm=float(fwhm),
        )
        
    except Exception as e:
        print(f"Simulation Error: {e}")


if __name__ == "__main__":
    main()

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

# Run the zone plate simulation
python zone_plate.py