"""
Metal Cavity LDOS Simulation with OptixLog Integration

Computes the local density of states (LDOS) in a metal cavity
and compares with Purcell formula 2Q/(πωV).

Based on the Meep tutorial: metal-cavity-ldos.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 metal_cavity(w, client=None, step_offset=0):
    """
    Compute LDOS for a metal cavity with aperture width w.
    """
    resolution = 50
    sxy = 2
    dpml = 1
    sxy_total = sxy + 2 * dpml
    cell = mp.Vector3(sxy_total, sxy_total)
    
    pml_layers = [mp.PML(dpml)]
    
    a = 1
    t = 0.1
    geometry = [
        mp.Block(mp.Vector3(a + 2 * t, a + 2 * t, mp.inf), material=mp.metal),
        mp.Block(mp.Vector3(a, a, mp.inf), material=mp.air),
    ]
    
    # Add aperture
    geometry.append(
        mp.Block(
            center=mp.Vector3(a / 2),
            size=mp.Vector3(2 * t, w, mp.inf),
            material=mp.air
        )
    )
    
    fcen = math.sqrt(0.5) / a
    df = 0.2
    
    sources = [
        mp.Source(
            src=mp.GaussianSource(fcen, fwidth=df),
            component=mp.Ez,
            center=mp.Vector3()
        )
    ]
    
    symmetries = [mp.Mirror(mp.Y)]
    
    sim = mp.Simulation(
        cell_size=cell,
        geometry=geometry,
        boundary_layers=pml_layers,
        sources=sources,
        symmetries=symmetries,
        resolution=resolution,
    )
    
    # Find resonance with Harminv
    h = mp.Harminv(mp.Ez, mp.Vector3(), fcen, df)
    sim.run(mp.after_sources(h), until_after_sources=500)
    
    m = h.modes[0]
    f = m.freq
    Q = m.Q
    Vmode = 0.25 * a * a
    ldos_formula = Q / Vmode / (2 * math.pi * f * math.pi * 0.5)
    
    sim.reset_meep()
    
    # Compute LDOS directly
    T = 2 * Q * (1 / f)
    sim.run(mp.dft_ldos(f, 0, 1), until_after_sources=T)
    ldos_meep = sim.ldos_data[0]
    
    return ldos_formula, ldos_meep, f, Q


def main():
    """Main simulation function for metal cavity LDOS."""
    
    if not optixlog.is_master_process():
        return
    
    try:
        client = optixlog.init(
            api_key=api_key,
            api_url=api_url,
            project=project_name,
            run_name="metal_cavity_ldos",
            config={
                "simulation_type": "ldos",
                "description": "Metal cavity LDOS analysis",
            },
            create_project_if_not_exists=True
        )
        
        client.log(
            step=0,
            resolution=50,
            cavity_size=1.0,
            metal_thickness=0.1,
        )
        
        # Aperture widths to simulate
        ws = np.arange(0.2, 0.5, 0.1)
        ldos_formula = np.zeros(len(ws))
        ldos_meep = np.zeros(len(ws))
        Q_factors = np.zeros(len(ws))
        
        for j, w in enumerate(ws):
            ldos_formula[j], ldos_meep[j], _, Q_factors[j] = metal_cavity(w)
            
            client.log(
                step=1 + j,
                aperture_width=float(w),
                Q_factor=float(Q_factors[j]),
                ldos_formula=float(ldos_formula[j]),
                ldos_meep=float(ldos_meep[j]),
            )
        
        client.log(
            step=100,
            simulation_completed=True,
            max_Q=float(np.max(Q_factors)),
            max_ldos=float(np.max(ldos_meep))
        )
        
    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 LDOS simulation
python metal-cavity-ldos.py