"""
Mie Scattering Simulation with OptixLog Integration

Computes the scattering cross-section of a dielectric sphere.
"""

import os
import meep as mp
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import numpy as np
from optixlog import Optixlog


def compute_mie_scattering(r, n_sphere, wvl, resolution=25):
    """Compute scattering cross-section for a single wavelength"""
    
    # Cell parameters
    dpml = 1.0
    dair = 2.0
    s = 2 * (r + dair + dpml)
    cell = mp.Vector3(s, s)
    
    pml_layers = [mp.PML(dpml)]
    
    # Frequency
    fcen = 1 / wvl
    
    # Sphere geometry
    geometry = [
        mp.Sphere(
            radius=r,
            center=mp.Vector3(),
            material=mp.Medium(index=n_sphere)
        )
    ]
    
    # Total-field/scattered-field source
    sources = [
        mp.Source(
            src=mp.GaussianSource(fcen, fwidth=0.2*fcen, is_integrated=True),
            center=mp.Vector3(-0.5*s + dpml + 0.5*dair),
            size=mp.Vector3(0, s),
            component=mp.Ez,
        )
    ]
    
    # Create simulation without sphere (for normalization)
    sim_empty = mp.Simulation(
        resolution=resolution,
        cell_size=cell,
        boundary_layers=pml_layers,
        sources=sources,
    )
    
    # Flux monitor enclosing sphere
    flux_box_empty = sim_empty.add_flux(
        fcen, 0, 1,
        mp.FluxRegion(center=mp.Vector3(-r-0.5*dair, 0), size=mp.Vector3(0, 2*r+dair)),
        mp.FluxRegion(center=mp.Vector3(r+0.5*dair, 0), size=mp.Vector3(0, 2*r+dair), weight=-1),
        mp.FluxRegion(center=mp.Vector3(0, r+0.5*dair), size=mp.Vector3(2*r+dair, 0), weight=-1),
        mp.FluxRegion(center=mp.Vector3(0, -r-0.5*dair), size=mp.Vector3(2*r+dair, 0)),
    )
    
    sim_empty.run(until_after_sources=50)
    empty_flux = mp.get_fluxes(flux_box_empty)[0]
    
    # Create simulation with sphere
    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell,
        boundary_layers=pml_layers,
        sources=sources,
        geometry=geometry,
    )
    
    flux_box = sim.add_flux(
        fcen, 0, 1,
        mp.FluxRegion(center=mp.Vector3(-r-0.5*dair, 0), size=mp.Vector3(0, 2*r+dair)),
        mp.FluxRegion(center=mp.Vector3(r+0.5*dair, 0), size=mp.Vector3(0, 2*r+dair), weight=-1),
        mp.FluxRegion(center=mp.Vector3(0, r+0.5*dair), size=mp.Vector3(2*r+dair, 0), weight=-1),
        mp.FluxRegion(center=mp.Vector3(0, -r-0.5*dair), size=mp.Vector3(2*r+dair, 0)),
    )
    
    sim.run(until_after_sources=50)
    scattered_flux = empty_flux - mp.get_fluxes(flux_box)[0]
    
    # Scattering cross-section (normalized by incident intensity)
    sigma_sca = scattered_flux / (empty_flux / s)
    
    # Scattering efficiency
    Q_sca = sigma_sca / (np.pi * r**2)
    
    return Q_sca, sigma_sca


def main():
    api_key = os.getenv("OPTIX_API_KEY", "your_api_key_here")
    
    client = Optixlog(api_key=api_key)
    project = client.project(name="MeepExamples", create_if_not_exists=True)
    
    # Parameters
    r = 1.0        # sphere radius
    n_sphere = 2.0  # refractive index
    
    # Wavelength sweep
    wvl_min = 0.5
    wvl_max = 3.0
    n_wvl = 20
    wavelengths = np.linspace(wvl_min, wvl_max, n_wvl)
    
    run = project.run(
        name="mie_scattering_analysis",
        config={
            "simulation_type": "mie_scattering",
            "sphere_radius": r,
            "refractive_index": n_sphere,
            "wavelength_min": wvl_min,
            "wavelength_max": wvl_max,
            "num_wavelengths": n_wvl
        }
    )
    
    run.log(step=0, sphere_radius=r, refractive_index=n_sphere)
    
    # Compute scattering for each wavelength
    Q_sca_list = []
    size_params = []
    
    print("⚡ Computing scattering cross-sections...")
    for i, wvl in enumerate(wavelengths):
        x = 2 * np.pi * r / wvl  # Size parameter
        size_params.append(x)
        
        Q_sca, sigma_sca = compute_mie_scattering(r, n_sphere, wvl)
        Q_sca_list.append(Q_sca)
        
        if i % 5 == 0:
            print(f"   λ = {wvl:.2f}, x = {x:.2f}, Q_sca = {Q_sca:.3f}")
        
        run.log(step=i+1,
                wavelength=float(wvl),
                size_parameter=float(x),
                scattering_efficiency=float(Q_sca))
    
    # Create plots
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # Q_sca vs wavelength
    axes[0].plot(wavelengths, Q_sca_list, 'bo-', markersize=6)
    axes[0].set_xlabel('Wavelength (μm)')
    axes[0].set_ylabel('Scattering Efficiency Q_sca')
    axes[0].set_title(f'Mie Scattering (r={r} μm, n={n_sphere})')
    axes[0].grid(True, alpha=0.3)
    
    # Q_sca vs size parameter
    axes[1].plot(size_params, Q_sca_list, 'ro-', markersize=6)
    axes[1].set_xlabel('Size Parameter x = 2πr/λ')
    axes[1].set_ylabel('Scattering Efficiency Q_sca')
    axes[1].set_title('Scattering vs Size Parameter')
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    run.log_matplotlib("mie_scattering_spectrum", fig)
    plt.close(fig)
    
    # Summary
    max_Q = max(Q_sca_list)
    max_idx = Q_sca_list.index(max_Q)
    
    run.log(step=n_wvl + 1,
            max_Q_sca=float(max_Q),
            max_Q_wavelength=float(wavelengths[max_idx]),
            max_Q_size_param=float(size_params[max_idx]),
            simulation_completed=True)
    
    print(f"\n📊 Maximum Q_sca: {max_Q:.3f} at λ={wavelengths[max_idx]:.2f} μm")
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

geometry = [
    mp.Sphere(
        radius=r,
        center=mp.Vector3(),
        material=mp.Medium(index=n_sphere)
    )
]

# Empty simulation flux (incident)
empty_flux = mp.get_fluxes(flux_box_empty)[0]

# With sphere (transmitted)
sphere_flux = mp.get_fluxes(flux_box)[0]

# Scattered = incident - transmitted
scattered_flux = empty_flux - sphere_flux

Q_sca = sigma_sca / (np.pi * r**2)