"""
Faraday Rotation Simulation with OptixLog Integration

Simulates Faraday rotation of a linearly polarized plane wave
through a gyrotropic medium with an applied magnetic field bias.

Based on the Meep tutorial: faraday-rotation.py
"""

import os
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 main():
    """Main simulation function for Faraday rotation."""
    
    if not optixlog.is_master_process():
        return
    
    try:
        client = optixlog.init(
            api_key=api_key,
            api_url=api_url,
            project=project_name,
            run_name="faraday_rotation_simulation",
            config={
                "simulation_type": "faraday_rotation",
                "description": "Faraday rotation in gyrotropic medium",
            },
            create_project_if_not_exists=True
        )
        
        # Material parameters
        epsn = 1.5     # background permittivity
        f0 = 1.0       # natural frequency
        gamma = 1e-6   # damping rate
        sn = 0.1       # sigma parameter
        b0 = 0.15      # magnitude of bias vector
        
        client.log(
            step=0,
            resolution=50,
            background_permittivity=epsn,
            natural_frequency=f0,
            bias_magnitude=b0
        )
        
        # Define gyrotropic material
        susc = [
            mp.GyrotropicLorentzianSusceptibility(
                frequency=f0,
                gamma=gamma,
                sigma=sn,
                bias=mp.Vector3(0, 0, b0)
            )
        ]
        mat = mp.Medium(epsilon=epsn, mu=1, E_susceptibilities=susc)
        
        # Simulation parameters
        tmax = 100
        L = 20.0
        cell = mp.Vector3(0, 0, L)
        fsrc = 0.8
        src_z = -8.5
        resolution = 50
        
        pml_layers = [mp.PML(thickness=1.0, direction=mp.Z)]
        
        sources = [
            mp.Source(
                mp.ContinuousSource(frequency=fsrc),
                component=mp.Ex,
                center=mp.Vector3(0, 0, src_z),
            )
        ]
        
        sim = mp.Simulation(
            cell_size=cell,
            geometry=[],
            sources=sources,
            boundary_layers=pml_layers,
            default_material=mat,
            resolution=resolution,
        )
        
        sim.run(until=tmax)
        
        client.log(step=1, simulation_completed=True)
        
        # Get field data
        ex_data = sim.get_efield_x().real
        ey_data = sim.get_efield_y().real
        z = np.linspace(-L / 2, L / 2, len(ex_data))
        
        # Calculate analytic result
        dfsq = f0**2 - 1j * fsrc * gamma - fsrc**2
        eperp = epsn + sn * f0**2 * dfsq / (dfsq**2 - (fsrc * b0) ** 2)
        eta = sn * f0**2 * fsrc * b0 / (dfsq**2 - (fsrc * b0) ** 2)
        
        k_gyro = 2 * np.pi * fsrc * np.sqrt(0.5 * (eperp - np.sqrt(eperp**2 - eta**2)))
        
        client.log(
            step=2,
            max_ex=float(np.max(np.abs(ex_data))),
            max_ey=float(np.max(np.abs(ey_data))),
            k_gyro_real=float(k_gyro.real),
        )
        
        # Calculate Verdet constant
        verdet_constant = k_gyro.real / (2 * np.pi)
        
        client.log(
            step=4,
            verdet_constant=float(verdet_constant.real),
            rotation_rate_deg_per_um=float(np.degrees(k_gyro.real)),
        )
        
    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 Faraday rotation simulation
python faraday-rotation.py

ε = | ε_perp   -iη      0    |
    | iη       ε_perp   0    |
    | 0        0        ε_z  |