"""
Gaussian Beam Simulation with OptixLog Integration

Simulates a focused Gaussian beam propagating through space.
"""

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


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
    s = 14  # cell size
    resolution = 50
    dpml = 2
    
    beam_x0 = mp.Vector3(0, 3.0)  # beam focus position
    rot_angle = 0  # rotation angle (degrees)
    beam_kdir = mp.Vector3(0, 1, 0).rotate(
        mp.Vector3(0, 0, 1), math.radians(rot_angle)
    )
    beam_w0 = 0.8  # beam waist
    beam_E0 = mp.Vector3(0, 0, 1)  # polarization (Ez)
    fcen = 1  # frequency
    
    run = project.run(
        name="gaussian_beam_simulation",
        config={
            "simulation_type": "gaussian_beam",
            "resolution": resolution,
            "cell_size": s,
            "beam_waist": beam_w0,
            "frequency": fcen,
            "rotation_angle": rot_angle,
            "focus_y": beam_x0.y
        }
    )
    
    run.log(step=0, beam_waist=beam_w0, frequency=fcen)
    
    # Derived parameters
    wavelength = 1.0 / fcen
    rayleigh_range = np.pi * beam_w0**2 / wavelength
    divergence = wavelength / (np.pi * beam_w0)
    
    print(f"Beam parameters:")
    print(f"  Waist: {beam_w0}")
    print(f"  Rayleigh range: {rayleigh_range:.2f}")
    print(f"  Divergence: {np.degrees(divergence):.2f}°")
    
    # Setup simulation
    cell_size = mp.Vector3(s, s)
    boundary_layers = [mp.PML(thickness=dpml)]
    
    sources = [
        mp.GaussianBeamSource(
            src=mp.ContinuousSource(fcen),
            center=mp.Vector3(0, -0.5 * s + dpml + 1.0),
            size=mp.Vector3(s),
            beam_x0=beam_x0,
            beam_kdir=beam_kdir,
            beam_w0=beam_w0,
            beam_E0=beam_E0,
        )
    ]
    
    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell_size,
        boundary_layers=boundary_layers,
        sources=sources,
    )
    
    print("⚡ Running simulation...")
    sim.run(until=20)
    
    # Get field data
    ez_data = sim.get_array(
        center=mp.Vector3(), 
        size=mp.Vector3(s - 2*dpml, s - 2*dpml), 
        component=mp.Ez
    )
    
    field_magnitude = np.abs(ez_data)
    max_field = np.max(field_magnitude)
    
    run.log(step=1, max_field_amplitude=float(max_field),
            rayleigh_range=rayleigh_range)
    
    # Create visualization
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # Field pattern
    output_plane = mp.Volume(
        center=mp.Vector3(), 
        size=mp.Vector3(s - 2*dpml, s - 2*dpml)
    )
    plt.sca(axes[0])
    sim.plot2D(fields=mp.Ez, output_plane=output_plane, ax=axes[0])
    axes[0].set_title(f'Gaussian Beam Ez Field (w₀={beam_w0})')
    
    # Intensity profile
    im = axes[1].imshow(
        field_magnitude.T**2,
        cmap='hot',
        origin='lower',
        aspect='auto'
    )
    axes[1].set_title('Beam Intensity Profile')
    axes[1].set_xlabel('X position')
    axes[1].set_ylabel('Y position')
    plt.colorbar(im, ax=axes[1], label='|Ez|²')
    
    plt.tight_layout()
    run.log_matplotlib("gaussian_beam_analysis", fig)
    plt.close(fig)
    
    # Cross-section analysis
    fig2, axes2 = plt.subplots(1, 2, figsize=(12, 5))
    
    # Y cross-section (along propagation)
    x_mid = ez_data.shape[0] // 2
    y_profile = field_magnitude[x_mid, :]**2
    axes2[0].plot(y_profile, 'b-', linewidth=2)
    axes2[0].set_xlabel('Y position')
    axes2[0].set_ylabel('Intensity')
    axes2[0].set_title('Intensity Along Propagation')
    axes2[0].grid(True, alpha=0.3)
    
    # X cross-section at focus
    y_focus = int(beam_x0.y * resolution / (s - 2*dpml) * ez_data.shape[1] / 2 
                  + ez_data.shape[1] / 2)
    y_focus = min(y_focus, ez_data.shape[1] - 1)
    x_profile = field_magnitude[:, y_focus]**2
    x_coords = np.linspace(-0.5*(s-2*dpml), 0.5*(s-2*dpml), len(x_profile))
    
    # Fit Gaussian
    from scipy.optimize import curve_fit
    try:
        def gaussian(x, A, x0, sigma):
            return A * np.exp(-(x - x0)**2 / (2 * sigma**2))
        popt, _ = curve_fit(gaussian, x_coords, x_profile, 
                           p0=[max(x_profile), 0, beam_w0])
        measured_waist = abs(popt[2]) * np.sqrt(2)  # Convert sigma to 1/e² waist
        
        axes2[1].plot(x_coords, x_profile, 'b-', linewidth=2, label='Meep')
        axes2[1].plot(x_coords, gaussian(x_coords, *popt), 'r--', 
                     linewidth=2, label=f'Fit (w={measured_waist:.3f})')
        axes2[1].legend()
        
        run.log(step=2, measured_waist=float(measured_waist),
                input_waist=beam_w0)
    except:
        axes2[1].plot(x_coords, x_profile, 'b-', linewidth=2)
    
    axes2[1].set_xlabel('X position')
    axes2[1].set_ylabel('Intensity')
    axes2[1].set_title('Transverse Profile at Focus')
    axes2[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    run.log_matplotlib("beam_cross_sections", fig2)
    plt.close(fig2)
    
    run.log(step=3, analysis_complete=True)
    
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

mp.GaussianBeamSource(
    src=mp.ContinuousSource(fcen),
    center=mp.Vector3(0, y_source),  # Source plane location
    size=mp.Vector3(width),          # Source width
    beam_x0=mp.Vector3(0, y_focus),  # Focus position
    beam_kdir=mp.Vector3(0, 1, 0),   # Propagation direction
    beam_w0=0.8,                      # Waist radius
    beam_E0=mp.Vector3(0, 0, 1),     # Polarization
)

beam_kdir = mp.Vector3(0, 1, 0).rotate(
    mp.Vector3(0, 0, 1),  # Rotation axis (z)
    math.radians(angle)    # Angle in radians
)