import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

import tidy3d as td
from tidy3d.plugins.mode import ModeSolver
from tidy3d import web

# Define structure related parameters
sim_size = [20, 20, 8.5]
sim_x = sim_size[0]
sim_y = sim_size[1]
sim_z = sim_size[2]
grid_spec = td.GridSpec(wavelength = 0.85)

lambda0  = 0.85
lambda_low = 0.75
lambda_high = 0.95
npoints = 200
freq0 = td.C_0 / lambda0
fwidth = td.C_0 * (1/lambda_low - 1/lambda_high)
freqs = np.linspace(freq0 - fwidth/2, freq0 + fwidth/2, npoints)
run_time = 20 / fwidth

"""
from types import SimpleNamespace
# GUI medium used in script need add here after pasting
param = SimpleNamespace()
"""

r = .498  # Radial increment
R0 = .645 # Initial inner radius
l = .224 # Arc length between holes
a = .075 # Hole radius
n_rings = 12 # Number of rings
t_hbn = .2 # hBN thickness
center = [0, 0] # Center of the hole array

# Define the media 
Vacuum = td.Medium(
    name = 'Vacuum', 
)

Glass = td.Medium(
    name = 'Glass', 
    permittivity = 2.25, 
)

mat_hbn = td.Medium.from_nk(
    n = 2.1, 
    k = 0, 
    freq = freq0, 
    name = 'hbn'
)


# Define the source, a single oscillating point dipole
pointdipole_0 = td.PointDipole(
    name = 'pointdipole_0', 
    source_time = td.GaussianPulse(freq0 = freq0, fwidth = fwidth), 
    polarization = 'Ey'
)


# Define the monitors
fluxmonitor_0 = td.FluxMonitor(
    name = 'fluxmonitor_0', 
    size = [0, sim_y, sim_z], 
    freqs = freqs, 
)

fluxmonitor_1 = td.FluxMonitor(
    name = 'fluxmonitor_1', 
    center = [0, 0, 3], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)

fluxmonitor_2 = td.FluxMonitor(
    name = 'fluxmonitor_2', 
    center = [0, 0, -3], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)

fluxmonitor_3 = td.FluxMonitor(
    name = 'fluxmonitor_3', 
    size = [sim_x, 0, sim_z], 
    freqs = freqs, 
)

fieldmonitor_4 = td.FieldMonitor(
    name = 'fieldmonitor_4', 
    size = [0, sim_y, sim_z], 
    freqs = freqs, 
)

fieldmonitor_5 = td.FieldMonitor(
    name = 'fieldmonitor_5', 
    size = [sim_x, 0, sim_z], 
    freqs = freqs, 
)

fieldmonitor_6 = td.FieldMonitor(
    name = 'fieldmonitor_6', 
    center = [0, 0, 3], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)

fieldmonitor_7 = td.FieldMonitor(
    name = 'fieldmonitor_7', 
    center = [0, 0, -3], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)

fieldmonitor_8 = td.FieldMonitor(
    name = 'fieldmonitor_8', 
    center = [0, 0, -2], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)

fieldmonitor_9 = td.FieldMonitor(
    name = 'fieldmonitor_9', 
    center = [0, 0, -1], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)

fieldmonitor_10 = td.FieldMonitor(
    name = 'fieldmonitor_10', 
    center = [0, 0, -0.2], 
    size = [sim_x, sim_y, 0], 
    freqs = freqs, 
)


# Define the structures

# Glass cylinder, representing optic fiber
cylinder_glass = td.Structure(
    geometry = td.Cylinder(radius = 50, center = [0, 0, -2.55], length = 4.9), 
    name = 'cylinder_glass', 
    medium = Glass
)

# hBN slab with no holes
cylinder_hbn = td.Structure(
    geometry = td.Cylinder(
        center = [0, 0, 0],
        radius = R0 + r*(n_rings+1),
        length = t_hbn,
        axis = 2
    ),
    medium = mat_hbn
)


# Function to generate the holes filled with air
def generate_object_0(center, r, R0, l, a, n_rings, t_hbn):
    hole_positions = []
    radii = np.linspace(R0, R0 + r * n_rings, n_rings)
    for R in radii:
        n_angles = int(2 * np.pi * R / l)
        for i in range(n_angles):
            angle = i * 2 * np.pi / n_angles 
            x = center[0] + R * np.cos(angle)
            y = center[1] + R * np.sin(angle)
            hole_positions.append([x, y])

    mat_air = td.Medium(permittivity = 1, name = 'air')
    mat_hbn = td.Medium(permittivity = 5.1, name = 'hbn')

    structures = []
    
    for pos in hole_positions:
        structures.append(
            td.Structure(
                geometry = td.Cylinder(
                    center = [pos[0], pos[1], 0],
                    radius = a,
                    length = t_hbn,
                    axis = 2
                ),
                medium = mat_air
            )
        )
        
    return structures


generate_object_0_result = generate_object_0(center, r, R0, l, a, n_rings, t_hbn)
generate_object_0_result = generate_object_0_result if isinstance(generate_object_0_result, list) else [generate_object_0_result]

def getObjectInArray(className, items):
    return list(filter(lambda item: isinstance(item, className), items))

# Define the simulation objects

sim_with_cbg = td.Simulation(
    size = sim_size, 
    grid_spec = grid_spec,
    run_time = {'source_factor' : 3, 'quality_factor' : 1},
    medium = Vacuum,
    sources = [pointdipole_0, *getObjectInArray(td.Source, [*generate_object_0_result])],
    monitors = [fluxmonitor_0, fluxmonitor_1, fluxmonitor_2, fluxmonitor_3, fieldmonitor_4, fieldmonitor_5, fieldmonitor_6, fieldmonitor_7, fieldmonitor_8, *getObjectInArray(td.Monitor, [*generate_object_0_result])],
    structures = [cylinder_hbn, *getObjectInArray(td.Structure, generate_object_0_result), cylinder_glass],
)

sim_without_cbg = sim_with_cbg.copy(update = {'structures' : [cylinder_hbn, cylinder_glass]})