"""
Antenna PEC Ground Plane Analysis with OptixLog Integration

Computes the radiation pattern of a dipole antenna above a perfect-electric
conductor (PEC) ground plane and compares with analytic theory.
"""

import os
import math
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)
    
    # Simulation parameters
    resolution = 200
    n = 1.2  # refractive index
    h = 1.25  # antenna height above ground
    wvl = 0.65  # wavelength
    r = 1000 * wvl  # far-field radius
    npts = 50  # number of angles
    src_cmpt = mp.Ez
    
    run = project.run(
        name="antenna_pec_ground_plane",
        config={
            "simulation_type": "antenna_pec_ground_plane",
            "resolution": resolution,
            "refractive_index": n,
            "antenna_height": h,
            "wavelength": wvl,
            "polarization": "Ez"
        }
    )
    
    run.log(step=0, refractive_index=n, antenna_height=h, wavelength=wvl)

    angles = 0.5 * math.pi / npts * np.arange(npts)

    def radial_flux(sim, nearfield_box, r):
        """Compute radial flux at far-field points"""
        E = np.zeros((npts, 3), dtype=np.complex128)
        H = np.zeros((npts, 3), dtype=np.complex128)

        for i in range(npts):
            ff = sim.get_farfield(
                nearfield_box,
                mp.Vector3(r * math.sin(angles[i]), r * math.cos(angles[i]))
            )
            E[i, :] = [np.conj(ff[j]) for j in range(3)]
            H[i, :] = [ff[j + 3] for j in range(3)]

        Px = np.real(E[:, 1] * H[:, 2] - E[:, 2] * H[:, 1])
        Py = np.real(E[:, 2] * H[:, 0] - E[:, 0] * H[:, 2])
        return np.sqrt(np.square(Px) + np.square(Py))

    def free_space_radiation(src_cmpt):
        """Simulate free-space radiation"""
        sxy = 4
        dpml = 1
        cell_size = mp.Vector3(sxy + 2 * dpml, sxy + 2 * dpml)
        pml_layers = [mp.PML(dpml)]
        fcen = 1 / wvl

        sources = [
            mp.Source(
                src=mp.GaussianSource(fcen, fwidth=0.2 * fcen),
                center=mp.Vector3(),
                component=src_cmpt,
            )
        ]

        if src_cmpt == mp.Hz:
            symmetries = [mp.Mirror(mp.X, phase=-1), mp.Mirror(mp.Y, phase=-1)]
        else:
            symmetries = [mp.Mirror(mp.X, phase=+1), mp.Mirror(mp.Y, phase=+1)]

        sim = mp.Simulation(
            cell_size=cell_size,
            resolution=resolution,
            sources=sources,
            symmetries=symmetries,
            boundary_layers=pml_layers,
            default_material=mp.Medium(index=n),
        )

        nearfield_box = sim.add_near2far(
            fcen, 0, 1,
            mp.Near2FarRegion(center=mp.Vector3(0, +0.5 * sxy), size=mp.Vector3(sxy, 0)),
            mp.Near2FarRegion(center=mp.Vector3(0, -0.5 * sxy), size=mp.Vector3(sxy, 0), weight=-1),
            mp.Near2FarRegion(center=mp.Vector3(+0.5 * sxy, 0), size=mp.Vector3(0, sxy)),
            mp.Near2FarRegion(center=mp.Vector3(-0.5 * sxy, 0), size=mp.Vector3(0, sxy), weight=-1),
        )

        sim.run(until_after_sources=mp.stop_when_dft_decayed())
        return radial_flux(sim, nearfield_box, r)

    def pec_ground_plane_radiation(src_cmpt):
        """Simulate radiation with PEC ground plane using image method"""
        L = 8.0
        dpml = 1.0
        sxy = dpml + L + dpml
        cell_size = mp.Vector3(sxy, sxy, 0)
        boundary_layers = [mp.PML(dpml)]
        fcen = 1 / wvl

        # Two antennas of opposite sign (image method)
        sources = [
            mp.Source(
                src=mp.GaussianSource(fcen, fwidth=0.2 * fcen),
                component=src_cmpt,
                center=mp.Vector3(0, +h),
            ),
            mp.Source(
                src=mp.GaussianSource(fcen, fwidth=0.2 * fcen),
                component=src_cmpt,
                center=mp.Vector3(0, -h),
                amplitude=-1,  # Image source with opposite phase
            ),
        ]

        if src_cmpt == mp.Hz:
            symmetries = [mp.Mirror(mp.X, phase=-1)]
        else:
            symmetries = [mp.Mirror(mp.X, phase=+1)]

        sim = mp.Simulation(
            cell_size=cell_size,
            resolution=resolution,
            sources=sources,
            symmetries=symmetries,
            boundary_layers=boundary_layers,
            default_material=mp.Medium(index=n),
        )

        nearfield_box = sim.add_near2far(
            fcen, 0, 1,
            mp.Near2FarRegion(center=mp.Vector3(0, h), size=mp.Vector3(4 * h, 0), weight=+1),
            mp.Near2FarRegion(center=mp.Vector3(0, -h), size=mp.Vector3(4 * h, 0), weight=-1),
            mp.Near2FarRegion(center=mp.Vector3(h, 0), size=mp.Vector3(0, 4 * h), weight=+1),
            mp.Near2FarRegion(center=mp.Vector3(-h, 0), size=mp.Vector3(0, 4 * h), weight=-1),
        )

        sim.run(until_after_sources=mp.stop_when_dft_decayed())
        return radial_flux(sim, nearfield_box, r)

    # Run simulations
    print("⚡ Running free-space simulation...")
    Pr_fsp = free_space_radiation(src_cmpt)
    run.log(step=1, phase="free_space_completed")
    
    print("⚡ Running PEC ground plane simulation...")
    Pr_pec = pec_ground_plane_radiation(src_cmpt)
    run.log(step=2, phase="pec_simulation_completed")

    # Compute theoretical pattern (array factor)
    k = 2 * np.pi / (wvl / n)
    Pr_theory = np.zeros(npts)
    for i, ang in enumerate(angles):
        Pr_theory[i] = Pr_fsp[i] * 2 * np.sin(k * h * np.cos(ang))

    # Normalize
    Pr_pec_norm = Pr_pec / np.max(Pr_pec)
    Pr_theory_norm = (Pr_theory / max(Pr_theory)) ** 2

    # Create plots
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # Linear plot
    axes[0].plot(np.degrees(angles), Pr_pec_norm, "b-", linewidth=2, label="Meep")
    axes[0].plot(np.degrees(angles), Pr_theory_norm, "r--", linewidth=2, label="Theory")
    axes[0].set_xlabel("Angle (degrees)")
    axes[0].set_ylabel("Normalized Radial Flux")
    axes[0].set_title(f"Antenna above PEC Ground Plane (h={h} μm)")
    axes[0].set_xlim(0, 90)
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Error plot
    error = np.abs(Pr_pec_norm - Pr_theory_norm)
    axes[1].plot(np.degrees(angles), error, "g-", linewidth=2)
    axes[1].set_xlabel("Angle (degrees)")
    axes[1].set_ylabel("Absolute Error")
    axes[1].set_title("Pattern Error")
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    run.log_matplotlib("radiation_pattern_comparison", fig)
    plt.close(fig)

    # Calculate metrics
    correlation = np.corrcoef(Pr_pec_norm, Pr_theory_norm)[0, 1]
    
    run.log(step=3,
            correlation_coefficient=correlation,
            wavevector=k,
            mean_error=float(np.mean(error)),
            simulation_completed=True)
    
    print(f"📊 Correlation: {correlation:.6f}")
    print(f"\n✅ Simulation complete!")


if __name__ == "__main__":
    main()

sources = [
    mp.Source(..., center=mp.Vector3(0, +h)),       # Real source
    mp.Source(..., center=mp.Vector3(0, -h), amplitude=-1),  # Image source
]

k = 2 * np.pi / (wvl / n)
for i, ang in enumerate(angles):
    Pr_theory[i] = Pr_fsp[i] * 2 * np.sin(k * h * np.cos(ang))