import math
import numpy as np
import tidy3d as td
import tidy3d.web as web
import matplotlib.pyplot as plt

print(f"Tidy3D version: {td.__version__}")

# Ring geometry
N_RINGS = 5
R = 20.0           # ring radius (um)
W = 1.4            # waveguide width (um)
G = 0.1            # coupling gap (um) = 100 nm
PITCH = 50.0       # ring-to-ring pitch (um)

# Rib waveguide cross-section
SLAB_H = 0.1       # slab thickness = 100 nm
RIB_H = 0.5        # rib height = 500 nm
TOTAL_LN = SLAB_H + RIB_H  # total LN = 600 nm

# Derived
R_outer = R + W / 2      # 20.7 um
R_inner = R - W / 2      # 19.3 um
R_coupling = R + W/2 + G + W/2  # 21.5 um (center-to-center)

# Zigzag angle
sin_theta = 2.0 * R_coupling / PITCH
THETA = math.asin(sin_theta)
cos_theta = math.cos(THETA)
print(f"Zigzag angle: {math.degrees(THETA):.1f} deg")
print(f"R_coupling = {R_coupling} um")

# Ring center positions
RING_CENTERS = [(i * PITCH, 0.0) for i in range(N_RINGS)]

# Bus waveguide extensions
BUS_EXTEND = 4.0
BUS_EXTEND_IO = 8.0

# Wavelength / frequency
WL_CENTER = 1.55   # um
WL_SPAN = 0.02     # 1540-1560 nm
WL_MIN = WL_CENTER - WL_SPAN / 2
WL_MAX = WL_CENTER + WL_SPAN / 2
FREQ_CENTER = td.C_0 / WL_CENTER
FREQ_MIN = td.C_0 / WL_MAX
FREQ_MAX = td.C_0 / WL_MIN
FWIDTH = FREQ_MAX - FREQ_MIN

# Material index
N_LN = 2.211  # n_o for X-cut TFLN, TE polarization

# Z coordinates
Z_SLAB_TOP = SLAB_H   # 0.1
Z_RIB_TOP = TOTAL_LN  # 0.6

# Simulation settings
FDTD_Z_MIN = -1.0
FDTD_Z_MAX = 1.5
RUN_TIME_S = 60e-12   # 60 ps
FREQ_PTS = 300
STRUCT_MARGIN = 2.0

LN = td.Medium(permittivity=N_LN**2)
SiO2 = td.Medium(permittivity=1.44**2)

print(f"LiNbO3: n = {N_LN}, eps = {N_LN**2:.3f}")
print(f"SiO2: n = 1.44, eps = {1.44**2:.3f}")

def compute_buses():
    """Compute bus waveguide endpoints for zigzag cascade."""
    buses = []
    st, ct = sin_theta, cos_theta

    # Input bus (horizontal, tangent to ring 1 top)
    cx0, cy0 = RING_CENTERS[0]
    buses.append({
        'type': 'horizontal', 'is_input': True, 'is_output': False,
        'y': cy0 + R_coupling,
        'x_start': cx0 - BUS_EXTEND_IO, 'x_end': cx0 + BUS_EXTEND,
    })

    # Inter-ring diagonal buses
    for k in range(N_RINGS - 1):
        bus_idx = k + 1
        cx_from, cy_from = RING_CENTERS[k]
        cx_to, cy_to = RING_CENTERS[k + 1]

        if bus_idx % 2 == 1:  # backslash direction
            tp_from = (cx_from + R_coupling*st, cy_from + R_coupling*ct)
            tp_to = (cx_to - R_coupling*st, cy_to - R_coupling*ct)
        else:  # slash direction
            tp_from = (cx_from + R_coupling*st, cy_from - R_coupling*ct)
            tp_to = (cx_to - R_coupling*st, cy_to + R_coupling*ct)

        dx = tp_to[0] - tp_from[0]
        dy = tp_to[1] - tp_from[1]
        seg_len = math.sqrt(dx**2 + dy**2)
        ux, uy = dx/seg_len, dy/seg_len
        angle = math.atan2(dy, dx)

        buses.append({
            'type': 'diagonal', 'is_input': False, 'is_output': False,
            'angle': angle,
            'p_start': (tp_from[0] - BUS_EXTEND*ux, tp_from[1] - BUS_EXTEND*uy),
            'p_end': (tp_to[0] + BUS_EXTEND*ux, tp_to[1] + BUS_EXTEND*uy),
        })

    # Output bus (horizontal, tangent to ring 5 top)
    cx_last, cy_last = RING_CENTERS[-1]
    buses.append({
        'type': 'horizontal', 'is_input': False, 'is_output': True,
        'y': cy_last + R_coupling,
        'x_start': cx_last - BUS_EXTEND, 'x_end': cx_last + BUS_EXTEND_IO,
    })

    return buses

buses = compute_buses()
print(f"Total buses: {len(buses)} (1 input + {N_RINGS-1} diagonal + 1 output)")

structures = []

# Collect extents for domain sizing
all_x, all_y = [], []
for bus in buses:
    if bus['type'] == 'horizontal':
        all_x += [bus['x_start'], bus['x_end']]
        all_y += [bus['y']]
    else:
        all_x += [bus['p_start'][0], bus['p_end'][0]]
        all_y += [bus['p_start'][1], bus['p_end'][1]]
for cx, cy in RING_CENTERS:
    all_x += [cx - R_outer, cx + R_outer]
    all_y += [cy - R_outer, cy + R_outer]

sub_x_c = (min(all_x) + max(all_x)) / 2
sub_y_c = (min(all_y) + max(all_y)) / 2
sub_x_span = max(all_x) - min(all_x) + 6
sub_y_span = max(all_y) - min(all_y) + 6

# 1. SiO2 substrate
structures.append(td.Structure(
    geometry=td.Box(center=[sub_x_c, sub_y_c, -1.0], size=[sub_x_span, sub_y_span, 2.0]),
    medium=SiO2, name="substrate",
))

# 2. LN slab (continuous)
structures.append(td.Structure(
    geometry=td.Box(center=[sub_x_c, sub_y_c, SLAB_H/2], size=[sub_x_span, sub_y_span, SLAB_H]),
    medium=LN, name="LN_slab",
))

# 3. Ring ridges
for i, (cx, cy) in enumerate(RING_CENTERS):
    outer_cyl = td.Cylinder(center=[cx, cy, Z_SLAB_TOP + RIB_H/2], axis=2, radius=R_outer, length=RIB_H)
    inner_cyl = td.Cylinder(center=[cx, cy, Z_SLAB_TOP + RIB_H/2], axis=2, radius=R_inner, length=RIB_H)
    structures.append(td.Structure(
        geometry=td.ClipOperation(operation="difference", geometry_a=outer_cyl, geometry_b=inner_cyl),
        medium=LN, name=f"ring_{i+1}",
    ))

# 4. Bus waveguide ridges
for i, bus in enumerate(buses):
    z_c = Z_SLAB_TOP + RIB_H / 2
    if bus['type'] == 'horizontal':
        label = "input" if bus['is_input'] else "output"
        xc = (bus['x_start'] + bus['x_end']) / 2
        xs = bus['x_end'] - bus['x_start']
        structures.append(td.Structure(
            geometry=td.Box(center=[xc, bus['y'], z_c], size=[xs, W, RIB_H]),
            medium=LN, name=f"bus_{label}",
        ))
    else:
        px_s, py_s = bus['p_start']
        px_e, py_e = bus['p_end']
        angle = bus['angle']
        dx = math.cos(angle + math.pi/2) * W/2
        dy = math.sin(angle + math.pi/2) * W/2
        vertices = [
            (px_s - dx, py_s - dy), (px_e - dx, py_e - dy),
            (px_e + dx, py_e + dy), (px_s + dx, py_s + dy),
        ]
        structures.append(td.Structure(
            geometry=td.PolySlab(vertices=vertices, slab_bounds=(Z_SLAB_TOP, Z_RIB_TOP), axis=2),
            medium=LN, name=f"bus_inter_{i}",
        ))

print(f"Total structures: {len(structures)}")

bus_in = buses[0]
src_x = bus_in['x_start'] + 3.0
src_y = bus_in['y']
src_z = 0.3
mon_y_span = W + 2.0
mon_z_span = 2.5

mode_spec = td.ModeSpec(num_modes=1, target_neff=N_LN)

# Broadband mode source
mode_source = td.ModeSource(
    center=[src_x, src_y, src_z],
    size=[0, mon_y_span, mon_z_span],
    source_time=td.GaussianPulse(freq0=FREQ_CENTER, fwidth=FWIDTH),
    direction="+", mode_spec=mode_spec, mode_index=0, num_freqs=7,
    name="mode_source",
)

# Frequency points for monitors
freqs_monitor = np.linspace(FREQ_MIN, FREQ_MAX, FREQ_PTS)

# Through port monitor
through_monitor = td.ModeMonitor(
    center=[src_x + 2.0, src_y, src_z],
    size=[0, mon_y_span, mon_z_span],
    freqs=freqs_monitor, mode_spec=mode_spec, name="through",
)

# Output (cascade drop) port monitor
bus_out = buses[-1]
output_monitor = td.ModeMonitor(
    center=[bus_out['x_end'] - 3.0, bus_out['y'], src_z],
    size=[0, mon_y_span, mon_z_span],
    freqs=freqs_monitor, mode_spec=mode_spec, name="output",
)

print(f"Source: ({src_x:.1f}, {src_y:.1f}) um")
print(f"Through monitor: ({src_x+2:.1f}, {src_y:.1f}) um")
print(f"Output monitor: ({bus_out['x_end']-3:.1f}, {bus_out['y']:.1f}) um")

fdtd_x_min = min(all_x) - STRUCT_MARGIN
fdtd_x_max = max(all_x) + STRUCT_MARGIN
fdtd_y_min = min(all_y) - STRUCT_MARGIN
fdtd_y_max = max(all_y) + STRUCT_MARGIN

sim = td.Simulation(
    center=[
        (fdtd_x_min + fdtd_x_max) / 2,
        (fdtd_y_min + fdtd_y_max) / 2,
        (FDTD_Z_MIN + FDTD_Z_MAX) / 2,
    ],
    size=[
        fdtd_x_max - fdtd_x_min,
        fdtd_y_max - fdtd_y_min,
        FDTD_Z_MAX - FDTD_Z_MIN,
    ],
    grid_spec=td.GridSpec.auto(min_steps_per_wvl=10, wavelength=WL_CENTER),
    structures=structures,
    sources=[mode_source],
    monitors=[through_monitor, output_monitor],
    run_time=RUN_TIME_S,
    boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
    medium=td.Medium(permittivity=1.0),  # air cladding
)

print(f"Domain: {sim.size[0]:.1f} x {sim.size[1]:.1f} x {sim.size[2]:.1f} um")
print(f"Grid cells: {sim.num_cells/1e6:.1f}M")
print(f"Time steps: {sim.num_time_steps:,}")
print(f"Run time: {RUN_TIME_S*1e12:.0f} ps")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Top view (XY plane through rib center)
sim.plot(z=0.35, ax=axes[0])
axes[0].set_title('Top view (z = 0.35 um)', fontsize=11)

# Cross-section (YZ plane through ring 1)
sim.plot(x=0.0, ax=axes[1])
axes[1].set_title('Cross-section at Ring 1 (x = 0)', fontsize=11)

plt.tight_layout()
plt.show()

import time

job = web.Job(simulation=sim, task_name="TFLN_5ring_cascade_MRR")

t0 = time.time()
sim_data = job.run(path="tidy3d_5ring_cascade.hdf5")
elapsed = time.time() - t0

print(f"Completed in {elapsed:.1f} s ({elapsed/60:.1f} min)")
print(f"Real cost: {web.real_cost(job.task_id):.2f} FlexCredits")

# Extract transmission spectra
freqs = sim_data["output"].amps.coords['f'].values
wavelengths = td.C_0 / freqs * 1e3  # nm

T_through = np.abs(sim_data["through"].amps.sel(mode_index=0, direction="+").values)**2
T_output = np.abs(sim_data["output"].amps.sel(mode_index=0, direction="+").values)**2

# Effective index
n_eff = np.real(sim_data["through"].n_eff.sel(mode_index=0).values)
print(f"Mode n_eff: {np.mean(n_eff):.4f}")
print(f"Through port max: {np.max(T_through):.4f}")
print(f"Output port max: {np.max(T_output):.6f}")

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

# (a) Transmission in dB
ax1.plot(wavelengths, 10*np.log10(np.clip(T_through, 1e-12, None)),
         'r-', lw=1.2, label='Through port')
ax1.plot(wavelengths, 10*np.log10(np.clip(T_output, 1e-12, None)),
         'b-', lw=1.2, label='Cascade output (5-ring drop)')
ax1.set_xlabel('Wavelength (nm)')
ax1.set_ylabel('Transmission (dB)')
ax1.set_title('5-Ring Cascade MRR in TFLN -- Transmission Spectrum')
ax1.legend()
ax1.grid(True, alpha=0.2)
ax1.set_xlim(1540, 1560)

# (b) Through port linear (zoom on resonances)
ax2.plot(wavelengths, T_through, 'r-', lw=1.2)
ax2.set_xlabel('Wavelength (nm)')
ax2.set_ylabel('Through port transmission')
ax2.set_title('Through Port (linear scale)')
ax2.grid(True, alpha=0.2)
ax2.set_xlim(1540, 1560)

plt.tight_layout()
plt.show()

from scipy.signal import find_peaks

# Find resonance dips in through port
thr_dB = 10 * np.log10(np.clip(T_through, 1e-12, None))
dips, props = find_peaks(-thr_dB, prominence=2, distance=10)

print(f"Found {len(dips)} resonance dips:")
for d in dips:
    wl_res = wavelengths[d]
    depth_dB = thr_dB[d]
    
    # Estimate FWHM
    half_depth = (0 + depth_dB) / 2  # half-way in dB
    left = d
    while left > 0 and thr_dB[left] < half_depth:
        left -= 1
    right = d
    while right < len(thr_dB)-1 and thr_dB[right] < half_depth:
        right += 1
    fwhm = abs(wavelengths[right] - wavelengths[left])  # nm
    Q = wl_res / fwhm if fwhm > 0 else 0
    
    print(f"  lambda = {wl_res:.3f} nm, depth = {depth_dB:.1f} dB, FWHM = {fwhm:.4f} nm, Q_L ~ {Q:.0f}")