"""
Stochastic Emitter Simulation with OptixLog Integration

Computes emission from stochastic dipole sources above a
metal/dielectric structure.
"""

import os
import argparse
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():
    parser = argparse.ArgumentParser()
    parser.add_argument("-res", type=int, default=50)
    parser.add_argument("-nr", type=int, default=20)
    parser.add_argument("-nd", type=int, default=10)
    parser.add_argument("-nf", type=int, default=500)
    parser.add_argument("-method", type=int, choices=[1, 2], default=1)
    args = parser.parse_args()
    
    if not optixlog.is_master_process():
        return
    
    try:
        from meep.materials import Ag
        
        client = optixlog.init(
            api_key=api_key,
            project=project_name,
            run_name=f"stochastic_emitter_method{args.method}",
            config={"simulation_type": "stochastic_emission"},
            create_project_if_not_exists=True
        )
        
        resolution = args.res
        dpml = 1.0
        dair = 1.0
        dsub = 5.0
        dAg = 0.5
        sx = 1.1
        sy = dpml + dair + dsub + dAg
        
        fcen = 1.0
        df = 0.2
        nfreq = args.nf
        ndipole = args.nd
        
        client.log(step=0, resolution=resolution, num_dipoles=ndipole, method=args.method)
        
        cell_size = mp.Vector3(sx, sy)
        pml_layers = [mp.PML(direction=mp.Y, thickness=dpml, side=mp.High)]
        
        geometry = [
            mp.Block(material=mp.Medium(index=3.45), center=mp.Vector3(0, 0.5 * sy - dpml - dair - 0.5 * dsub), size=mp.Vector3(mp.inf, dsub, mp.inf)),
            mp.Block(material=Ag, center=mp.Vector3(0, -0.5 * sy + 0.5 * dAg), size=mp.Vector3(mp.inf, dAg, mp.inf)),
        ]
        
        # Run simulations and collect flux data
        # (Implementation depends on method choice)
        
        client.log(step=100, simulation_completed=True)
        
    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()

# Method 1: Random dipole ensemble
python stochastic_emitter.py -method 1 -nr 20

# Method 2: Single dipole sweep
python stochastic_emitter.py -method 2 -nd 10