"""
1D Absorption Simulation with OptixLog Integration

Simulates electromagnetic field absorption in a 1D aluminum medium using Meep
and logs the results to OptixLog for visualization and tracking.

Usage:
    export OPTIX_API_KEY="your_api_key_here"
    python absorbed_1d.py                    # Without PML (uses absorber)
    python absorbed_1d.py --pml             # With PML boundary layers
"""

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):
    # Initialize OptixLog client
    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"absorbed_1d_aluminum_{boundary_type.lower()}",
        config={
            "simulation_type": "absorbed_1d",
            "material": "Aluminum",
            "dimensions": 1,
            "resolution": 40,
            "use_pml": args.pml,
            "boundary_type": boundary_type,
            "wavelength": 0.803,
            "cell_size_z": 10
        }
    )
    
    print(f"🚀 Starting 1D Absorption simulation")
    print(f"📊 Boundary type: {boundary_type}")

    # 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,
        )
    ]

    # Log simulation parameters
    run.log(step=0,
            resolution=resolution,
            cell_size_z=10,
            wavelength=wavelength,
            frequency=frequency,
            boundary_type=boundary_type,
            material="Aluminum",
            source_component="Ex")

    # Store field data
    field_data = []
    time_data = []
    step_counter = [0]

    def monitor_field(sim):
        p = sim.get_field_point(mp.Ex, mp.Vector3())
        current_time = sim.meep_time()
        
        field_data.append(p.real)
        time_data.append(current_time)
        
        # Log field data periodically
        if len(field_data) % 20 == 0:
            step_counter[0] += 1
            run.log(step=step_counter[0],
                    field_ex=p.real,
                    time=current_time,
                    field_magnitude=abs(p.real))

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

    print(f"⚡ Running simulation with {boundary_type} boundary...")
    
    sim.run(
        mp.at_every(10, monitor_field),
        until_after_sources=mp.stop_when_fields_decayed(50, mp.Ex, mp.Vector3(), 1e-6),
    )

    # Analyze and log final results
    if field_data:
        max_field = max(field_data)
        min_field = min(field_data)
        final_field = field_data[-1]
        field_decay = abs(final_field / max_field) if max_field != 0 else 0
        
        run.log(step=step_counter[0] + 1,
                max_field=max_field,
                min_field=min_field,
                final_field=final_field,
                field_decay_ratio=field_decay,
                total_time_steps=len(field_data),
                simulation_completed=True)

        # Create field evolution plot
        fig, axes = plt.subplots(2, 1, figsize=(12, 8))
        
        axes[0].plot(time_data, field_data, 'b-', linewidth=1.5, alpha=0.8)
        axes[0].set_xlabel('Time', fontsize=11)
        axes[0].set_ylabel('Ex Field (real part)', fontsize=11)
        axes[0].set_title(f'Field Evolution in 1D Aluminum Medium (Boundary: {boundary_type})', fontsize=12)
        axes[0].grid(True, alpha=0.3)
        
        # Log scale for decay analysis
        field_magnitude = [abs(f) for f in field_data]
        axes[1].semilogy(time_data, field_magnitude, 'r-', linewidth=1.5, alpha=0.8)
        axes[1].set_xlabel('Time', fontsize=11)
        axes[1].set_ylabel('|Ex Field|', fontsize=11)
        axes[1].set_title('Field Magnitude (Log Scale)', fontsize=12)
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        run.log_matplotlib("field_evolution", fig)
        plt.close(fig)

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


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pml", action="store_true", default=False,
                        help="Use PML boundary condition instead of absorber")
    args = parser.parse_args()
    main(args)

from meep.materials import Al  # Aluminum

sim = mp.Simulation(
    ...
    default_material=Al,
    ...
)

def monitor_field(sim):
    p = sim.get_field_point(mp.Ex, mp.Vector3())
    # p is a complex number: p.real, p.imag

# Call at regular intervals
sim.run(mp.at_every(10, monitor_field), until=100)

sim.run(
    until_after_sources=mp.stop_when_fields_decayed(
        50,          # dT: time interval for checking
        mp.Ex,       # component to monitor
        mp.Vector3(), # monitoring point
        1e-6         # decay threshold
    )
)

# Run with absorber boundary (default)
python absorbed_1d.py

# Run with PML boundary
python absorbed_1d.py --pml