"""
1D Absorber Simulation with OptixLog Integration

Compares absorber vs PML boundary conditions in a 1D simulation.
"""

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


def main(args):
    """Main simulation function with OptixLog integration"""
    
    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)
    
    boundary_type = "PML" if args.pml else "Absorber"
    
    run = project.run(
        name=f"absorber_1d_aluminum_{boundary_type.lower()}",
        config={
            "simulation_type": "absorber_1d",
            "material": "Aluminum",
            "dimensions": 1,
            "resolution": 40,
            "use_pml": args.pml,
            "boundary_type": boundary_type,
            "wavelength": 0.803,
            "cell_size_z": 10
        }
    )
    
    # Simulation parameters
    resolution = 40
    cell_size = mp.Vector3(z=10)
    wavelength = 0.803
    frequency = 1 / wavelength

    boundary_layers = [
        mp.PML(1, direction=mp.Z) if args.pml else mp.Absorber(1, direction=mp.Z)
    ]

    sources = [
        mp.Source(
            src=mp.GaussianSource(frequency, fwidth=0.1),
            center=mp.Vector3(),
            component=mp.Ex,
        )
    ]

    # Data collection
    times = []
    field_values = []
    field_amplitudes = []

    def monitor_field(sim):
        p = sim.get_field_point(mp.Ex, mp.Vector3())
        current_time = sim.meep_time()
        
        times.append(current_time)
        field_values.append(p.real)
        field_amplitudes.append(abs(p))
        
        print(f"ex:, {current_time}, {p.real}, {p.imag}, {abs(p)}")

    # Create and run simulation
    sim = mp.Simulation(
        cell_size=cell_size,
        resolution=resolution,
        dimensions=1,
        default_material=Al,
        boundary_layers=boundary_layers,
        sources=sources,
    )

    sim.run(
        mp.at_every(10, monitor_field),
        until_after_sources=mp.stop_when_fields_decayed(50, mp.Ex, mp.Vector3(), 1e-6),
    )

    # Create comprehensive analysis plots
    if times and field_values:
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        
        # Real part of field
        axes[0, 0].plot(times, field_values, 'b-', linewidth=1.5)
        axes[0, 0].set_xlabel('Time')
        axes[0, 0].set_ylabel('Ex Field (Real)')
        axes[0, 0].set_title('Electric Field Evolution')
        axes[0, 0].grid(True, alpha=0.3)
        
        # Amplitude
        axes[0, 1].plot(times, field_amplitudes, 'r-', linewidth=1.5)
        axes[0, 1].set_xlabel('Time')
        axes[0, 1].set_ylabel('|Ex| Amplitude')
        axes[0, 1].set_title('Field Amplitude')
        axes[0, 1].grid(True, alpha=0.3)
        
        # Decay analysis (log scale)
        if len(field_amplitudes) > 10:
            tail_start = len(field_amplitudes) // 2
            tail_times = np.array(times[tail_start:])
            tail_amps = np.array(field_amplitudes[tail_start:])
            
            axes[1, 0].semilogy(tail_times, tail_amps, 'go-', markersize=3)
            axes[1, 0].set_xlabel('Time')
            axes[1, 0].set_ylabel('|Ex| (log scale)')
            axes[1, 0].set_title('Field Decay Analysis')
            axes[1, 0].grid(True, alpha=0.3)
        
        # Final field profile
        z_coords = np.linspace(-5, 5, 100)
        field_profile = [abs(sim.get_field_point(mp.Ex, mp.Vector3(z=z))) for z in z_coords]
        
        axes[1, 1].plot(z_coords, field_profile, 'purple', linewidth=2)
        axes[1, 1].set_xlabel('Z Position')
        axes[1, 1].set_ylabel('|Ex| Amplitude')
        axes[1, 1].set_title('Final Field Distribution')
        axes[1, 1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        run.log_matplotlib("field_analysis", fig)
        plt.close(fig)

        # Calculate and log statistics
        max_field = max(field_amplitudes)
        final_field = field_amplitudes[-1]
        decay_ratio = final_field / max_field if max_field > 0 else 0
        
        run.log(step=2,
                max_field_amplitude=max_field,
                final_field_amplitude=final_field,
                field_decay_ratio=decay_ratio,
                simulation_time=times[-1],
                total_time_steps=len(times))

    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pml", action="store_true", default=False)
    args = parser.parse_args()
    main(args)

z_coords = np.linspace(-5, 5, 100)
field_profile = [
    abs(sim.get_field_point(mp.Ex, mp.Vector3(z=z))) 
    for z in z_coords
]

# With absorber boundary
python absorber-1d.py

# With PML boundary
python absorber-1d.py --pml