"""
Cherenkov Radiation Simulation with OptixLog Integration

Simulates a moving point charge with superluminal phase velocity
in a dielectric medium emitting Cherenkov radiation.
"""

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 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
    sx = 60
    sy = 60
    resolution = 10
    dpml = 1.0
    v = 0.7  # velocity (units of c)
    n = 1.5  # refractive index
    
    # Check Cherenkov condition
    beta_n = v * n
    if beta_n > 1:
        cherenkov_angle = np.arccos(1 / beta_n)
        cherenkov_angle_deg = np.degrees(cherenkov_angle)
    else:
        cherenkov_angle_deg = 0
    
    run = project.run(
        name="cherenkov_radiation_simulation",
        config={
            "simulation_type": "cherenkov_radiation",
            "resolution": resolution,
            "cell_size": sx,
            "charge_velocity": v,
            "medium_index": n,
            "beta_n": beta_n,
            "cherenkov_angle": cherenkov_angle_deg,
            "superluminal": beta_n > 1
        }
    )
    
    run.log(step=0, charge_velocity=v, medium_index=n, beta_n=beta_n,
            cherenkov_angle=cherenkov_angle_deg)
    
    print(f"βn = {beta_n:.2f} {'> 1 (radiation expected)' if beta_n > 1 else '< 1'}")
    if beta_n > 1:
        print(f"Theoretical Cherenkov angle: {cherenkov_angle_deg:.1f}°")
    
    # Setup simulation
    cell_size = mp.Vector3(sx, sy, 0)
    pml_layers = [mp.PML(thickness=dpml)]
    symmetries = [mp.Mirror(direction=mp.Y)]
    
    sim = mp.Simulation(
        resolution=resolution,
        cell_size=cell_size,
        default_material=mp.Medium(index=n),
        symmetries=symmetries,
        boundary_layers=pml_layers,
    )
    
    # Track field snapshots
    field_snapshots = []
    time_points = []
    
    def move_source(sim_obj):
        """Move the point source according to velocity."""
        x_pos = -0.5 * sx + dpml + v * sim_obj.meep_time()
        if x_pos < 0.5 * sx - dpml:
            sim_obj.change_sources([
                mp.Source(
                    mp.ContinuousSource(frequency=1e-10),
                    component=mp.Ex,
                    center=mp.Vector3(x_pos),
                )
            ])
    
    # Data collection
    snapshot_interval = 5
    last_snapshot = [0]
    
    def collect_data(sim_obj):
        current_time = sim_obj.meep_time()
        if current_time - last_snapshot[0] >= snapshot_interval:
            hz_data = sim_obj.get_array(
                center=mp.Vector3(), 
                size=mp.Vector3(sx-2*dpml, sy-2*dpml), 
                component=mp.Hz
            )
            field_snapshots.append(hz_data.copy())
            time_points.append(current_time)
            last_snapshot[0] = current_time
    
    print("⚡ Running simulation...")
    total_time = sx / v
    
    sim.run(move_source, collect_data, until=total_time)
    
    run.log(step=1, simulation_completed=True, total_time=total_time)
    
    # Analyze and visualize
    if field_snapshots:
        final_hz = sim.get_array(
            center=mp.Vector3(), 
            size=mp.Vector3(sx-2*dpml, sy-2*dpml), 
            component=mp.Hz
        )
        
        max_field = np.max(np.abs(final_hz))
        
        run.log(step=2, max_field_amplitude=float(max_field),
                num_snapshots=len(field_snapshots))
        
        # Create visualization
        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
        
        extent = [-0.5*(sx-2*dpml), 0.5*(sx-2*dpml), 
                  -0.5*(sy-2*dpml), 0.5*(sy-2*dpml)]
        
        # Hz field pattern
        im1 = axes[0].imshow(
            final_hz.T,
            cmap='RdBu',
            origin='lower',
            extent=extent,
            aspect='equal'
        )
        axes[0].set_title(f'Cherenkov Radiation Pattern\nθ_c = {cherenkov_angle_deg:.1f}°')
        axes[0].set_xlabel('X position')
        axes[0].set_ylabel('Y position')
        plt.colorbar(im1, ax=axes[0], label='Hz')
        
        # Field evolution
        if len(field_snapshots) > 1:
            max_fields = [np.max(np.abs(fs)) for fs in field_snapshots]
            axes[1].plot(time_points, max_fields, 'b-', linewidth=2)
            axes[1].set_xlabel('Simulation Time')
            axes[1].set_ylabel('Max Field Amplitude')
            axes[1].set_title('Field Evolution')
            axes[1].grid(True, alpha=0.3)
        else:
            intensity = np.abs(final_hz)**2
            im2 = axes[1].imshow(intensity.T, cmap='hot', origin='lower',
                                  extent=extent, aspect='equal')
            axes[1].set_title('Field Intensity')
            plt.colorbar(im2, ax=axes[1], label='|Hz|²')
        
        plt.tight_layout()
        run.log_matplotlib("cherenkov_radiation_analysis", fig)
        plt.close(fig)
    
    run.log(step=3, theoretical_angle=cherenkov_angle_deg,
            analysis_complete=True)
    
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

def move_source(sim_obj):
    x_pos = -0.5 * sx + dpml + v * sim_obj.meep_time()
    sim_obj.change_sources([
        mp.Source(
            mp.ContinuousSource(frequency=1e-10),
            component=mp.Ex,
            center=mp.Vector3(x_pos),
        )
    ])

beta_n = v * n  # Effective velocity in medium
if beta_n > 1:
    # Cherenkov radiation occurs
    angle = np.arccos(1 / beta_n)