import tidy3d as td
import tidy3d.web as web
import numpy as np
import matplotlib.pyplot as plt
import autograd.numpy as anp
import autograd as ag
import optax
from tidy3d.plugins.autograd import make_filter_and_project, rescale, make_erosion_dilation_penalty, value_and_grad, make_curvature_penalty
from tidy3d.plugins.mode import ModeSolver
from tidy3d.plugins.mode.web import run as run_mode_solver

# 1. GENERAL SETUP

print(" Step 1: General Setup ")

wavelength = 1.0
freq0 = td.C_0 / wavelength

# Materials
eps_wg = 2.75
wg_medium = td.Medium(permittivity=eps_wg)

# Dimensions
lx_conv = 5.0  # Converter Length
ly_conv = 3.0  # Converter Width
wg_width_conv = 1.2
output_wg_length = 3.0

# Grid & Design Resolution
dl_design = 0.01
nx = int(lx_conv / dl_design)
ny = int(ly_conv / dl_design)

# Initial Random Parameters
# Seed the RNG so the optimization starts from a reproducible initial density.
np.random.seed(0)
params0_conv = np.random.random((nx, ny))

print(f"Setup Complete. Wavelength: {wavelength} um")

# STEP 1.5: MODE ANALYSIS

print(" Step 1.5: Running Mode Analysis ")

wg_input_analysis = td.Structure(
    geometry=td.Box(center=(-lx_conv, 0, 0), size=(lx_conv, wg_width_conv, td.inf)), # center = (-5.0, 0, 0), size = (5.0, 1.2, inf)
    medium=wg_medium
)

mode_size = (0, wg_width_conv * 4, td.inf) # size = (0, 4.8, inf)
source_x_pos = -lx_conv # -5.0

plane = td.Box(center=[source_x_pos, 0, 0], size=mode_size) # size = (0, 4.8, inf)

sim_mode = td.Simulation(
    size=(1, ly_conv + 2, 0), # size = (1, 5.0, 0)
    center=(source_x_pos, 0, 0), # center = (-5.0, 0, 0)
    grid_spec=td.GridSpec.auto(min_steps_per_wvl=20, wavelength=wavelength), # min_steps_per_wvl=20
    structures=[wg_input_analysis],
    run_time=1e-12,
    boundary_spec=td.BoundarySpec.pml(x=False, y=True, z=False),
    sources=[],
    monitors=[]
)

num_modes = 4
mode_spec_analysis = td.ModeSpec(num_modes=num_modes)

mode_solver = ModeSolver(
    simulation=sim_mode,
    plane=plane,
    mode_spec=mode_spec_analysis,
    freqs=[freq0],
)

print("Running Advanced Mode Analysis on Server...")
modes = run_mode_solver(mode_solver, reduce_simulation=True, verbose=True)

print("Plotting Fields...")
fig, axs = plt.subplots(num_modes, 3, figsize=(12, 3 * num_modes), tight_layout=True)
for mode_index in range(num_modes):
    vmax = 1.1 * max(abs(modes.field_components[n].sel(mode_index=mode_index)).max() for n in ("Ex", "Ey", "Ez"))
    for field_name, ax in zip(("Ex", "Ey", "Ez"), axs[mode_index]):
        field = modes.field_components[field_name].sel(mode_index=mode_index)
        field.real.plot(label="Real", ax=ax)
        field.imag.plot(ls="--", label="Imag", ax=ax)
        ax.set_title(f"Mode {mode_index}, Field = {field_name}")
        ax.set_ylim(-vmax, vmax)
        ax.grid(True, alpha=0.3)
axs[0, 0].legend()
plt.show()
print("Effective indices: ", np.array(modes.n_eff))

# PART A: MODE CONVERTER OPTIMIZATION

print(" PART A: MODE CONVERTER OPTIMIZATION ")

radius = 0.2
filter_project = make_filter_and_project(radius, dl_design)

def get_eps(params, beta):
    processed_params = filter_project(params, beta)
    eps = rescale(processed_params, 1, eps_wg)
    return eps

# Visualization: Initial Structure
print("Plotting Initial Random Permittivity")
eps_initial = get_eps(params0_conv, beta=1).reshape((nx, ny))
plt.figure(figsize=(6, 4))
plt.imshow(np.flipud(eps_initial.T), cmap="gray", origin='lower',
           extent=[-lx_conv / 2, lx_conv / 2, -ly_conv / 2, ly_conv / 2])
plt.colorbar(label="Permittivity")
plt.title("Initial Random Structure")
plt.show()

# Optimization Loop (Mode Converter)
print("Starting Optimization Loop of Mode Converter...")
optimizer_conv = optax.adam(learning_rate=0.3)
params_conv = np.array(params0_conv)
opt_state_conv = optimizer_conv.init(params_conv)
dJ_conv_fn = value_and_grad(J_conv)

beta_history = []
steps_conv = 50

for i in range(steps_conv):
    current_beta = 1.0 + i * 2.0
    beta_history.append(current_beta)

    # Visualization: Iteration Steps
    eps_current = get_eps(params_conv, current_beta).reshape((nx, ny))
    plt.figure(figsize=(4, 2))
    plt.imshow(np.flipud(eps_current.T), cmap="gray", origin='lower')
    plt.title(f"Iteration {i + 1} (Beta={current_beta})")
    plt.axis("off")
    plt.show()

    val, grad = dJ_conv_fn(params_conv, beta=current_beta, step_num=i + 1, verbose=False)
    updates, opt_state_conv = optimizer_conv.update(-grad, opt_state_conv, params_conv)
    params_conv = optax.apply_updates(params_conv, updates)
    params_conv = np.clip(params_conv, 0, 1)
    print(f"Conv Step {i + 1}/{steps_conv} | Beta: {current_beta:.1f} | Obj: {val:.4f}")

# Final Verification
sim_final_conv = make_converter_sim(params_conv, beta=beta_history[-1])
field_mnt_conv = td.FieldMonitor(center=(0, 0, 0), size=(td.inf, td.inf, 0), freqs=[freq0], name="field_plot")
sim_final_conv = sim_final_conv.copy(update={"monitors": [field_mnt_conv] + list(sim_final_conv.monitors)})
data_final_conv = web.run(sim_final_conv, task_name="final_conv_verify", verbose=True)
eff_conv = measure_power_conv(data_final_conv)
print(f"Final Converter Efficiency (Mode 2): {eff_conv * 100:.2f}%")
sim_final_conv.to_gds_file("optimized_mode_converter.gds", z=0, permittivity_threshold=(1 + eps_wg) / 2,
                           frequency=freq0)

# PART B: TAPER OPTIMIZATION

print(" PART B: TAPER OPTIMIZATION ")

L_taper = 10.0
w_taper_in = wg_width_conv
w_taper_out = 5.0
buffer = 2

Lx_taper_sim = L_taper + 2 * buffer
Ly_taper_sim = w_taper_out + 2 * buffer

num_points = 15
xs = np.linspace(-L_taper / 2, L_taper / 2, num_points)

def get_ys(params):
    p_norm = (anp.tanh(np.pi * params) + 1) / 2
    widths = p_norm * (w_taper_out - w_taper_in) + w_taper_in
    return widths / 2.0

def get_params_from_ys(ys):
    full_widths = ys * 2.0
    p_norm = (full_widths - w_taper_in) / (w_taper_out - w_taper_in)
    p_norm = np.clip(p_norm, 1e-5, 1 - 1e-5)
    return np.arctanh(2 * p_norm - 1) / np.pi

# Optimization Loop (Taper)
print("Initializing Taper Optimization...")
ys_init = np.linspace(w_taper_in / 2, w_taper_out / 2, num_points)
params_taper = get_params_from_ys(ys_init)

optimizer_taper = optax.adam(learning_rate=0.02)
opt_state_taper = optimizer_taper.init(params_taper)
grad_fn_taper = ag.value_and_grad(J_taper)

for i in range(2):
    params_np = np.array(params_taper)
    val, grad = grad_fn_taper(params_np)
    grad = np.nan_to_num(grad)
    updates, opt_state_taper = optimizer_taper.update(-grad, opt_state_taper, params_taper)
    params_taper = optax.apply_updates(params_taper, updates)
    print(f"Taper Step {i + 1}/2 | Obj: {val:.4f}")

# Final Verification (Taper)
sim_final_taper = make_taper_sim(params_taper)
field_mnt_taper = td.FieldMonitor(center=(0, 0, 0), size=(td.inf, td.inf, 0), freqs=[freq0], name="field_plot")
sim_final_taper = sim_final_taper.copy(update={"monitors": [field_mnt_taper] + list(sim_final_taper.monitors)})
data_final_taper = web.run(sim_final_taper, task_name="final_taper_verify", verbose=True)
eff_taper = float(
    np.sum(np.abs(data_final_taper["taper_measure"].amps.sel(direction="+", f=freq0, mode_index=2).values) ** 2))

print(f"FINAL RESULTS:")
print(f"1. Converter Efficiency (Mode 0 -> 2): {eff_conv * 100:.2f}%")
print(f"2. Taper Efficiency (Mode 2 -> 2):     {eff_taper * 100:.2f}%")

sim_final_taper.to_gds_file("optimized_taper_mode2.gds", z=0, permittivity_threshold=(1 + eps_wg) / 2, frequency=freq0)

# PART C: FINAL COMBINED SIMULATION (Converter + INTERMEDIATE WG + Taper)

print(" PART C: FINAL COMBINED VERIFICATION ")

mode_spec_final = td.ModeSpec(num_modes=5)

# 1. Converter
# Converter stays at Left: Center = -2.5, Length=5
# Ends at x=0
eps_data_final = get_eps(params_conv, beta=beta_history[-1]).reshape((nx, ny, 1))
converter_geo_final = td.Box(center=(-lx_conv / 2, 0, 0), size=(lx_conv, ly_conv, td.inf))
struct_conv_final = td.Structure.from_permittivity_array(geometry=converter_geo_final, eps_data=eps_data_final)

#  2. Intermediate Waveguide
# Starts at x=0, Ends at x=2 (Length = 2.0 um)
L_mid = 2.0
wg_mid_geo = td.Box(
    center=(L_mid / 2, 0, 0),
    size=(L_mid, wg_width_conv, td.inf)  # Width matches Converter Output & Taper Input
)
struct_wg_mid = td.Structure(geometry=wg_mid_geo, medium=wg_medium)

#  3. Taper
# Previously started at x=0. Now must start at x=L_mid (2.0)
ys_opt = get_ys(params_taper)
xs_shifted = xs + L_taper / 2 + L_mid  # Shift x coordinates by L_taper/2 (to zero) + L_mid

vertices_final = anp.concatenate([
    anp.column_stack((xs_shifted, ys_opt)),
    anp.column_stack((xs_shifted[::-1], -ys_opt[::-1]))
])
taper_geo_final = td.PolySlab(vertices=vertices_final, slab_bounds=(-td.inf, td.inf), axis=2)
struct_taper_final = td.Structure(geometry=taper_geo_final, medium=wg_medium)

#  4. Input & Output Waveguides
wg_in_final = td.Structure(geometry=td.Box(center=(-lx_conv - 1.0, 0, 0), size=(2.0, wg_width_conv, td.inf)),
                           medium=wg_medium)
# Output WG starts at L_mid + L_taper
wg_out_final = td.Structure(geometry=td.Box(center=(L_mid + L_taper + 1.0, 0, 0), size=(2.0, w_taper_out, td.inf)),
                            medium=wg_medium)

#  5. Simulation Setup
Lx_total = lx_conv + L_mid + L_taper + 4.0
sim_center_x = (-lx_conv + L_mid + L_taper) / 2

sim_full = td.Simulation(
    size=(Lx_total, w_taper_out + 4.0, 0),
    center=(sim_center_x, 0, 0),
    grid_spec=td.GridSpec.auto(min_steps_per_wvl=25, wavelength=wavelength),
    structures=[wg_in_final, struct_conv_final, struct_wg_mid, struct_taper_final, wg_out_final],
    sources=[td.ModeSource(center=(-lx_conv - 0.5, 0, 0), size=(0, ly_conv + 2, td.inf),
                           source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10), direction='+', mode_index=0,
                           mode_spec=mode_spec_final)],
    monitors=[td.ModeMonitor(center=(L_mid + L_taper + 0.5, 0, 0), size=(0, w_taper_out + 2, td.inf), freqs=[freq0],
                             mode_spec=mode_spec_final, name="full_measure"),
              td.FieldMonitor(center=(sim_center_x, 0, 0), size=(td.inf, td.inf, 0), freqs=[freq0], name="field_plot")],
    run_time=200 / (freq0 / 10),
    boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=False)
)

#  VISUALIZATION of Final Device
print("Plotting Final Full Device Structure")
fig, ax = plt.subplots(figsize=(15, 6))
sim_full.plot_eps(z=0, ax=ax)
plt.title("Final Full Device Structure")
plt.xlabel("x (um)")
plt.show()

print("Running Full Combined System...")
data_full = web.run(sim_full, task_name="full_device_simulation", verbose=True)

amp_full = data_full["full_measure"].amps.sel(direction="+", f=freq0, mode_index=2).values
eff_full = float(np.sum(np.abs(amp_full) ** 2))

print(f"TOTAL SYSTEM EFFICIENCY: {eff_full * 100:.2f}%")

fig, ax = plt.subplots(figsize=(15, 6))
data_full.plot_field("field_plot", "Ez", z=0, ax=ax)
plt.title(f"Combined Device Field (Eff: {eff_full * 100:.2f}%)")
plt.xlabel("x (um)")
plt.show()

sim_full.to_gds_file("Final_Full_Device_Mode0_to_Mode1_to_Taper.gds", z=0, permittivity_threshold=(1 + eps_wg) / 2, frequency=freq0)
print("Saved 'Final_Full_Device_Mode0_to_Mode1_to_Taper.gds'")