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

r = 0.498  # Radial increment
R0 = 0.645  # Initial inner radius
l = 0.224  # Arc length between holes
a = 0.075  # Hole radius
n_rings = 12  # Number of rings
t_hbn = 0.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]}
)