"""
Cavity Far-Field Analysis with OptixLog Integration

Computes the far-field radiation pattern of a photonic crystal cavity
using the near-to-far field transformation.
"""

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
    resolution = 20
    fcen = 0.25
    df = 0.2
    eps = 13
    w = 1.2
    r = 0.36
    d = 1.4
    N = 3
    dpad = 32
    dpml = 0.5 / (fcen - 0.5 * df)
    
    # Cell size
    sx = 2 * (dpad + dpml + N) + d - 1
    d1 = 0.2
    d2 = 2.0
    sy = w + 2 * (d1 + d2 + dpml)
    
    run = project.run(
        name="cavity_farfield_analysis",
        config={
            "simulation_type": "cavity_farfield",
            "resolution": resolution,
            "center_frequency": fcen,
            "waveguide_epsilon": eps,
            "waveguide_width": w,
            "hole_radius": r,
            "num_holes": N
        }
    )
    
    run.log(step=0, resolution=resolution, fcen=fcen, waveguide_width=w)

    cell = mp.Vector3(sx, sy, 0)

    # Geometry: waveguide with holes
    geometry = [
        mp.Block(
            center=mp.Vector3(),
            size=mp.Vector3(mp.inf, w, mp.inf),
            material=mp.Medium(epsilon=eps),
        )
    ]

    # Add periodic holes
    for i in range(N):
        geometry.append(mp.Cylinder(r, center=mp.Vector3(d / 2 + i)))
        geometry.append(mp.Cylinder(r, center=mp.Vector3(d / -2 - i)))

    pml_layers = [mp.PML(dpml)]

    sources = [
        mp.Source(
            src=mp.GaussianSource(fcen, fwidth=df),
            component=mp.Hz,
            center=mp.Vector3()
        )
    ]

    symmetries = [mp.Mirror(mp.X, phase=-1), mp.Mirror(mp.Y, phase=-1)]

    sim = mp.Simulation(
        cell_size=cell,
        geometry=geometry,
        sources=sources,
        symmetries=symmetries,
        boundary_layers=pml_layers,
        resolution=resolution,
    )

    # Near-to-far transformation regions
    nearfield = sim.add_near2far(
        fcen, 0, 1,
        mp.Near2FarRegion(mp.Vector3(0, 0.5 * w + d1), size=mp.Vector3(sx - 2 * dpml)),
        mp.Near2FarRegion(
            mp.Vector3(-0.5 * sx + dpml, 0.5 * w + 0.5 * d1),
            size=mp.Vector3(0, d1),
            weight=-1.0,
        ),
        mp.Near2FarRegion(
            mp.Vector3(0.5 * sx - dpml, 0.5 * w + 0.5 * d1),
            size=mp.Vector3(0, d1)
        ),
    )

    # DFT field monitor for comparison
    mon = sim.add_dft_fields(
        [mp.Hz],
        fcen, 0, 1,
        center=mp.Vector3(0, 0.5 * w + d1 + d2),
        size=mp.Vector3(sx - 2 * (dpad + dpml), 0),
    )

    print("⚡ Running simulation...")
    sim.run(until_after_sources=mp.stop_when_dft_decayed())

    # Get field data
    Hz_mon = sim.get_dft_array(mon, mp.Hz, 0)
    (x, y, z, w_arr) = sim.get_array_metadata(dft_cell=mon)

    # Compute far-field along same line
    ff = []
    for xc in x:
        ff_pt = sim.get_farfield(nearfield, mp.Vector3(xc, y[0]))
        ff.append(ff_pt[5])  # Hz component
    ff = np.array(ff)

    # Create comparison plots
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    # Real part
    axes[0].plot(x, np.real(Hz_mon), "bo-", markersize=4, label="DFT")
    axes[0].plot(x, np.real(ff), "ro-", markersize=4, label="N2F")
    axes[0].set_xlabel("x (μm)")
    axes[0].set_ylabel("Re(Hz)")
    axes[0].set_title("Real Part")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Imaginary part
    axes[1].plot(x, np.imag(Hz_mon), "bo-", markersize=4, label="DFT")
    axes[1].plot(x, np.imag(ff), "ro-", markersize=4, label="N2F")
    axes[1].set_xlabel("x (μm)")
    axes[1].set_ylabel("Im(Hz)")
    axes[1].set_title("Imaginary Part")
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    # Magnitude
    axes[2].plot(x, np.abs(Hz_mon), "bo-", markersize=4, label="DFT")
    axes[2].plot(x, np.abs(ff), "ro-", markersize=4, label="N2F")
    axes[2].set_xlabel("x (μm)")
    axes[2].set_ylabel("|Hz|")
    axes[2].set_title("Magnitude")
    axes[2].legend()
    axes[2].grid(True, alpha=0.3)
    
    plt.suptitle("Near2Far vs DFT Field Comparison")
    plt.tight_layout()
    
    run.log_matplotlib("field_comparison", fig)
    plt.close(fig)

    # Calculate correlation
    correlation_mag = np.corrcoef(np.abs(Hz_mon), np.abs(ff))[0, 1]
    
    run.log(step=2,
            correlation_magnitude=correlation_mag,
            max_field_dft=float(np.max(np.abs(Hz_mon))),
            max_field_n2f=float(np.max(np.abs(ff))),
            simulation_completed=True)
    
    print(f"📊 Magnitude correlation: {correlation_mag:.4f}")
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

nearfield = sim.add_near2far(
    fcen, 0, 1,
    mp.Near2FarRegion(center, size=mp.Vector3(length, 0)),
    mp.Near2FarRegion(center, size=mp.Vector3(0, length), weight=-1),  # Inward normal
    ...
)

mon = sim.add_dft_fields(
    [mp.Hz],
    fcen, 0, 1,
    center=mp.Vector3(0, y_far),
    size=mp.Vector3(length, 0),
)

ff_pt = sim.get_farfield(nearfield, mp.Vector3(x, y))
# Returns [Ex, Ey, Ez, Hx, Hy, Hz]
Hz = ff_pt[5]