import numpy as np
import matplotlib.pyplot as plt
import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.resonance import ResonanceFinder

import gdstk

lambda_min = 0.8
lambda_max = 1.1

monitor_lambda = np.linspace(lambda_min, lambda_max, 401)
monitor_freq = td.constants.C_0 / monitor_lambda

automesh_per_wvl = 40
sim_size = Lx, Ly, Lz = (5.5, 5.5, 2)
offset_monitor = 0.4

vacuum = td.Medium(permittivity = 1, name = 'vacuum')
GaAs_permittivity = 3.55**2
GaAs = td.Medium(permittivity = GaAs_permittivity, name = 'GaAs')

freq0  = 302.675e12
fwidth = 144.131e12

dipole_source = td.PointDipole(
    center = (0, 0, 0),
    source_time = td.GaussianPulse(
        freq0 = freq0,
        fwidth = fwidth
    ),
    polarization = 'Ex',
    name = 'dipole_source'
)

t_start = 10/fwidth
t_stop  = 30e-12

slab_side_length = 5
slab_height = 0.2
slab = td.Structure(
    geometry = td.Box(
        center = (0, 0, 0),
        size = (slab_side_length, slab_side_length, slab_height),
    ),
    medium = GaAs
)

lib = gdstk.Library()
cell = lib.new_cell("BULLSEYE")

central_inner_radius = 0.355 + 0.014
ring_width = 0.07 + np.array([-0.036, -0.015, +0.011, -0.009, -0.011, +0.012, +0.005])
ring_height = 0.2
ring_period = 0.1075

angle = 3
quadrants = 4
theta1 = angle * np.pi / 180
theta2 = (90 - angle) * np.pi / 180

ring_numbers  = np.size(ring_width)

x_origin = 0
y_origin = 0

inner_radius = central_inner_radius

for i in range(ring_numbers):

    outer_radius = inner_radius + ring_width[i]

    for j in range(quadrants):

      gc_slice = gdstk.ellipse(
          (x_origin, y_origin),
          inner_radius,
          outer_radius,
          theta1 + j*np.pi/2,
          theta2 + j*np.pi/2,
          layer = 1,
          datatype = 1,
          tolerance = 0.0001
      )
      cell.add(gc_slice)

    inner_radius = outer_radius + ring_period

gc_etch = td.Geometry.from_gds(
      cell,
      gds_layer = 1,
      axis = 2,
      slab_bounds = (-ring_height/2, ring_height/2)
  )

etch = td.Structure(
    geometry = gc_etch,
    medium = vacuum,
    name = "etch"
)

time_monitor_1 = td.FieldTimeMonitor(
    name = "time_monitor_1",
    size = [0.0, 0.0, 0.0],
    center = [0.1, 0.1, 0.0],
    start = t_start
)

time_monitor_2 = td.FieldTimeMonitor(
    name = "time_monitor_2",
    size = [0.0, 0.0, 0.0],
    center = [0.1, 0.2, 0.0],
    start = t_start
)

time_monitor_3 = td.FieldTimeMonitor(
    name = "time_monitor_3",
    size = [0.0, 0.0, 0.0],
    center = [0.2, 0.2, 0.0],
    start = t_start
)

time_monitor_4 = td.FieldTimeMonitor(
    name = "time_monitor_4",
    size = [0.0, 0.0, 0.0],
    center = [0.2, 0.1, 0.0],
    start = t_start
)

time_monitors = [time_monitor_1, time_monitor_2, time_monitor_3, time_monitor_4]

far_field_monitor = td.FieldProjectionAngleMonitor(
    center = (0.0, 0.0, offset_monitor),
    size = (td.inf, td.inf, 0),
    name = 'far_field_monitor',
    freqs = monitor_freq,
    normal_dir = '+',
    phi = np.linspace(0, 2*np.pi, 361),
    theta = np.linspace(0, np.pi, 181)
)

def gaussian_far_field_no_polarization_sin_debye_apod(theta, phi, NA, n=1.0):

    """
    Compute an ideal far-field Gaussian beam profile with Debye apodization,
    given the theta and phi angles, numerical aperture (NA), and refractive index (n).
    Parameters:
    phi, theta : 1D arrays (rad)
        Input angular coordinates
    NA, n : float 
        Numerical aperture and refractive index (of air)
    Returns:
    A : 2D array
        Ideal far-field amplitude profile with Debye apodization
    """
    theta_vals = np.asarray(theta)
    phi_vals   = np.asarray(phi)

    theta_grid, phi_grid = np.meshgrid(theta_vals, phi_vals, indexing='ij')

    # Gaussian term
    gauss = np.exp(-(np.sin(theta_grid) / NA)**2)

    # Debye apodization with safe clipping
    apod = np.sqrt(np.clip(np.cos(theta_grid), 0, None))

    return gauss * apod

def compute_gaussian_overlap_without_pol(
    phi, theta, vals, 
    NA = 0.68, 
    n = 1.0,
    theta_max_deg = None,
    gauss_func = gaussian_far_field_no_polarization_sin_debye_apod,
):
    """
    Compute the mode-overlap between the simulated far-field intensity (vals)
    and an ideal Gaussian far-field WITHOUT polarization effects.

    Overlap is computed as:
        | ∫ E_sim E_gauss dΩ|^2 / ( ∫|E_sim|^2 dΩ  ∫|E_gauss|^2 dΩ )

    Added FEATURE:
        If theta_max_deg is provided (in degrees), the integrals are evaluated
        ONLY for theta <= theta_max_deg.

    Parameters
    ----------
    phi, theta : 1D arrays (rad)
        Angular sampling coordinates of the simulated far-field.
    vals : 2D array
        Simulated intensity map, shape = (len(theta), len(phi)).
    NA, n : floats
        Gaussian beam NA and refractive index.
    theta_max_deg : float or None
        Upper limit of theta (in degrees) for restricting the overlap region.
        If None → use entire theta-domain.
    gauss_func : callable
        Function that generates the ideal Gaussian amplitude:
            gauss_func(theta, phi, NA, n) -> 2D array (len(theta), len(phi))

    Returns
    -------
    percentage_overlap : float
        Mode overlap (0–100%).
    """

    # Convert inputs to arrays
    phi   = np.asarray(phi)
    theta = np.asarray(theta)
    vals  = np.asarray(vals)

    assert vals.shape == (len(theta), len(phi)), \
        f"vals has shape {vals.shape}, but expected {(len(theta), len(phi))}"

    # --- Choose which Gaussian model to use via gauss_func ---
    E_gauss = gauss_func(theta, phi, NA, n)
    I_gauss = E_gauss**2

    # Sim amplitude = sqrt(intensity)
    I_sim = vals
    E_sim = np.sqrt(I_sim)

    # Meshgrid for solid-angle element
    theta_grid, phi_grid = np.meshgrid(theta, phi, indexing='ij')

    # Solid angle element (assume uniform grid)
    dtheta = float(np.mean(np.diff(theta)))
    dphi   = float(np.mean(np.diff(phi)))
    dOmega = np.sin(theta_grid) * dtheta * dphi

    # --------------------------------------
    #  Mask to restrict region if requested
    # --------------------------------------
    if theta_max_deg is not None:
        theta_max_rad = np.radians(theta_max_deg)
        mask = theta_grid <= theta_max_rad
    else:
        mask = np.ones_like(theta_grid, dtype=bool)

    # Apply mask
    E_sim_m   = E_sim[mask]
    E_gauss_m = E_gauss[mask]
    I_sim_m   = I_sim[mask]
    I_gauss_m = I_gauss[mask]
    dOmega_m  = dOmega[mask]

    # --------------------------
    #  Mode overlap numerator
    # --------------------------
    overlap = np.sum(E_sim_m * E_gauss_m * dOmega_m)
    numerator = np.abs(overlap)**2

    # --------------------------
    #  Normalization denominator
    # --------------------------
    norm_sim   = np.sum(I_sim_m   * dOmega_m)
    norm_gauss = np.sum(I_gauss_m * dOmega_m)
    denominator = norm_sim * norm_gauss

    eta = numerator / denominator if denominator > 0 else 0.0
    return float(100 * eta)

def compute_power_fraction_and_ratio(
    phi, theta, vals,
    NA_target,
    NA_gauss=0.68,
    n=1.0,
    theta_max_deg=None,
    gauss_func=gaussian_far_field_no_polarization_sin_debye_apod,
):
    r"""
    Compute the percentage of power captured within a target NA cone for:
      (1) simulated far-field intensity vals(θ,φ)
      (2) ideal Gaussian (amplitude) generated by gauss_func

    Then compute the ratio:
        ratio = (percent_sim / percent_gauss)

    Power is integrated over solid angle:
        P \propto \iint I(θ,φ) sinθ dθ dφ

    Optional FEATURE:
        If theta_max_deg is provided, both total-power integrals are computed
        only over θ <= theta_max_deg (same as your overlap style).

    Parameters
    ----------
    phi, theta : 1D arrays (rad)
    vals : 2D array
        Simulated intensity map, shape = (len(theta), len(phi)).
    NA_target : float
        NA defining the capture cone for the fraction calculation.
    NA_gauss : float
        NA used to *define the Gaussian model* in gauss_func (beam divergence).
    n : float
        Refractive index of surrounding medium for θ = asin(NA/n).
    theta_max_deg : float or None
        If provided, restrict all integrals to θ <= theta_max_deg.
    gauss_func : callable
        Gaussian amplitude model:
            gauss_func(theta, phi, NA_gauss, n) -> 2D amplitude array

    Returns
    -------
    percent_sim : float
        0–100 (%) power fraction within NA_target for simulated intensity.
    percent_gauss : float
        0–100 (%) power fraction within NA_target for Gaussian intensity |E|^2.
    ratio : float
        (percent_sim / percent_gauss); np.inf if percent_gauss == 0.
    """

    if gauss_func is None:
        raise ValueError("gauss_func must be provided (it should return Gaussian amplitude).")

    # Convert inputs to arrays
    phi   = np.asarray(phi)
    theta = np.asarray(theta)
    vals  = np.asarray(vals)

    assert vals.shape == (len(theta), len(phi)), \
        f"vals has shape {vals.shape}, but expected {(len(theta), len(phi))}"

    # --- Build Gaussian intensity from amplitude ---
    E_gauss = gauss_func(theta, phi, NA_gauss, n)
    I_gauss = np.abs(E_gauss)**2

    # Simulated intensity
    I_sim = vals

    # Meshgrid for solid-angle element (assume uniform grid like your function)
    theta_grid, phi_grid = np.meshgrid(theta, phi, indexing='ij')

    dtheta = float(np.mean(np.diff(theta)))
    dphi   = float(np.mean(np.diff(phi)))
    dOmega = np.sin(theta_grid) * dtheta * dphi

    # --------------------------------------
    # Masks: (A) optional theta restriction
    # --------------------------------------
    if theta_max_deg is not None:
        theta_max_rad = np.radians(theta_max_deg)
        mask_theta = theta_grid <= theta_max_rad
    else:
        mask_theta = np.ones_like(theta_grid, dtype=bool)

    # --------------------------------------
    # Masks: (B) NA capture cone restriction
    # --------------------------------------
    theta_cap = np.arcsin(np.clip(NA_target / float(n), -1.0, 1.0))
    mask_cap = theta_grid <= theta_cap

    # Combined masks
    mask_total = mask_theta
    mask_in    = mask_theta & mask_cap

    # --------------------------
    # Simulated power fractions
    # --------------------------
    P_sim_total = np.sum(I_sim[mask_total] * dOmega[mask_total])
    P_sim_in    = np.sum(I_sim[mask_in]    * dOmega[mask_in])

    frac_sim = (P_sim_in / P_sim_total) if P_sim_total > 0 else 0.0

    # --------------------------
    # Gaussian power fractions
    # --------------------------
    P_g_total = np.sum(I_gauss[mask_total] * dOmega[mask_total])
    P_g_in    = np.sum(I_gauss[mask_in]    * dOmega[mask_in])

    frac_gauss = (P_g_in / P_g_total) if P_g_total > 0 else 0.0

    percent_sim   = 100.0 * float(frac_sim)
    percent_gauss = 100.0 * float(frac_gauss)
    ratio = np.inf if percent_gauss == 0 else (percent_sim / percent_gauss)

    return percent_sim, percent_gauss, float(ratio)

def make_far_field_plot_without_pol(phi, theta, vals, NA = 0.68, n=1.0, 
                                    gauss_func = gaussian_far_field_no_polarization_sin_debye_apod,
                                    calculate_overlap=True):
    """
    Plots the far-field intensity distribution, optionally saving the figure and calculating the mode overlap and mentioning it on the plot.
    Parameters:
    theta, phi : 1D arrays (rad)
        Input angular coordinates
    vals : 2D array
        Far-field intensity values
    NA, n : float 
        Numerical aperture and refractive index (of air)
    calculate_overlap : bool
        Whether to calculate and display the mode overlap (Ideally when plotting the Gaussian beam, this should be False)
    """

    phi = np.asarray(phi)
    theta = np.asarray(theta)
    vals = np.asarray(vals)
    assert vals.shape == (len(theta), len(phi))

    if calculate_overlap:
        overlap = compute_gaussian_overlap_without_pol(
            phi, theta, vals,
            NA=NA, n=n,
            gauss_func=gauss_func
        )


    I = vals
    I_norm = I / (I.max() if I.max()>0 else 1.0)

    PHI, THETA = np.meshgrid(phi, theta, indexing = 'xy')

    theta_NA_deg = np.degrees(np.arcsin(min(NA/float(n), 1.0)))
    phi_circle = np.linspace(0, 2.0 * np.pi, 721)
    theta_circle = np.full_like(phi_circle, theta_NA_deg)

    fig1, ax1 = plt.subplots(subplot_kw={'projection': 'polar'}, figsize = (6.5,5.5))
    ax1.tick_params(axis='y', colors='white')
    pcm1 = ax1.pcolormesh(PHI, np.degrees(THETA), I_norm, shading='auto', cmap='magma')
    cbar1 = fig1.colorbar(pcm1, ax=ax1, label=r"$|E|^2$")
    ax1.set_ylim(0, 90)
    #ax1.set_title('Far_field: Normalized Intensity')
    ax1.plot(phi_circle, theta_circle, 'w--', lw = 2, label=f'NA={NA:.2f}, n = {n}')
    ax1.legend(loc='upper right')
    if calculate_overlap:
        ax1.text(0.02, 0.98, f'Mode overlap: {overlap:.2f}%', transform = ax1.transAxes, va = 'top', ha = 'left',
                bbox = dict(boxstyle = 'round, pad = 0.3', facecolor = 'black', alpha = 0.55), color = 'white')
        
    fig1.tight_layout()
    plt.show()
    return

f_index = 168
far_field_data = sim_data[far_field_monitor.name] 
Etheta = far_field_data.Etheta.isel(f = f_index, r=0).values
Ephi = far_field_data.Ephi.isel(f = f_index, r=0).values
phi_vals = far_field_data.phi  
theta_vals = far_field_data.theta

intensity = np.abs(Etheta)**2 + np.abs(Ephi)**2