"""
Multilevel Atom Laser Simulation with OptixLog Integration

Simulates laser dynamics in a multilevel atom gain medium
inside a cavity using MEEP's saturable gain model.
"""

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", "")
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="multilevel_atom_laser",
            config={"simulation_type": "laser_dynamics"},
            create_project_if_not_exists=True
        )
        
        # Simulation parameters
        resolution = 400
        ncav = 1.5
        Lcav = 1
        dpad = 1
        dpml = 1
        sz = Lcav + dpad + dpml
        
        # Gain parameters
        omega_a = 40
        freq_21 = omega_a / (2 * math.pi)
        gamma_perp = 4
        gamma_21 = (2 * gamma_perp) / (2 * math.pi)
        theta = 1
        sigma_21 = 2 * theta * theta * omega_a
        rate_21 = 0.005
        N0 = 37
        Rp = 0.0051
        
        client.log(step=0, resolution=resolution, emission_frequency=freq_21, pumping_rate=Rp)
        
        cell_size = mp.Vector3(z=sz)
        pml_layers = [mp.PML(dpml, side=mp.High)]
        
        transitions = [
            mp.Transition(1, 2, pumping_rate=Rp, frequency=freq_21, gamma=gamma_21, sigma_diag=mp.Vector3(sigma_21, 0, 0)),
            mp.Transition(2, 1, transition_rate=rate_21),
        ]
        
        ml_atom = mp.MultilevelAtom(sigma=1, transitions=transitions, initial_populations=[N0])
        two_level = mp.Medium(index=ncav, E_susceptibilities=[ml_atom])
        
        geometry = [
            mp.Block(center=mp.Vector3(z=-0.5 * sz + 0.5 * Lcav), size=mp.Vector3(mp.inf, mp.inf, Lcav), material=two_level)
        ]
        
        sim = mp.Simulation(cell_size=cell_size, resolution=resolution, boundary_layers=pml_layers, geometry=geometry, dimensions=1)
        sim.init_sim()
        
        def field_func(p):
            return 1 if p.z == -0.5 * sz + 0.5 * Lcav else 0
        sim.fields.initialize_field(mp.Ex, field_func)
        
        endt = 7000
        sim.run(until=endt)
        
        client.log(step=1, simulation_completed=True)
        
    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()