mode_data = web.run(mode_solver, "mode_solver_kerr")

f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(
    2, 3, tight_layout=True, figsize=(10, 6)
)

mode_data.Ex.sel(mode_index=0, f=freq0).abs.plot(ax=ax1)
mode_data.Ey.sel(mode_index=0, f=freq0).abs.plot(ax=ax2)
mode_data.Ez.sel(mode_index=0, f=freq0).abs.plot(ax=ax3)
mode_data.Ex.sel(mode_index=1, f=freq0).abs.plot(ax=ax4)
mode_data.Ey.sel(mode_index=1, f=freq0).abs.plot(ax=ax5)
mode_data.Ez.sel(mode_index=1, f=freq0).abs.plot(ax=ax6)
ax1.set_title("|Ex|: mode_index=0")
ax2.set_title("|Ey|: mode_index=0")
ax3.set_title("|Ez|: mode_index=0")
ax4.set_title("|Ex|: mode_index=1")
ax5.set_title("|Ey|: mode_index=1")
ax6.set_title("|Ez|: mode_index=1")

mode_data.to_dataframe()

# ---------------------------
# Monitor Settings
# ---------------------------
# Define measurement frequency range for the ring mode monitors
lambdas_measure = np.linspace(lambda_beg, lambda_end, 1001)
freqs_measure = td.C_0 / lambdas_measure[::-1]

# Create multiple ring mode monitors with different time windows using a loop
pulses = [[40, 50], [400, 410], [600, 610], [730, 740]]
ring_mode_monitors = []
for i, pulse in enumerate(pulses, start=1):
    apodization = td.ApodizationSpec(
        start=1e-12 * pulse[0], end=1e-12 * pulse[1], width=2e-13
    )
    monitor = td.ModeMonitor(
        size=ring_mode_plane.size,
        center=ring_mode_plane.center,
        freqs=freqs_measure,
        mode_spec=td.ModeSpec(num_modes=2),
        apodization=apodization,
        name=f"ring_mode_{i}",
    )
    ring_mode_monitors.append(monitor)

# Define output flux monitor
out_flux_monitor = td.FluxTimeMonitor(
    size=ring_mode_plane.size,
    center=ring_mode_plane.center,
    name="out_flux",
)

# ---------------------------
# Source Definition
# ---------------------------
mode_source = td.ModeSource(
    size=mode_plane.size,
    center=mode_plane.center,
    source_time=td.ContinuousWave(freq0=freq0, fwidth=fwidth, amplitude=amp),
    mode_spec=td.ModeSpec(num_modes=2),
    mode_index=1,
    direction="+",
    num_freqs=11,
)

# ---------------------------
# Overall Simulation Setup
# ---------------------------
sim = td.Simulation(
    normalize_index=None,
    center=[0, sim_center_y, 0],
    size=[x_span, y_span, 0],
    grid_spec=td.GridSpec.auto(
        min_steps_per_wvl=min_steps_per_wvl, wavelength=td.C_0 / f_middle
    ),
    structures=[waveguide, ring],
    sources=[mode_source],
    monitors=[out_flux_monitor] + ring_mode_monitors,
    run_time=run_time,
    boundary_spec=td.BoundarySpec(
        x=td.Boundary.pml(), y=td.Boundary.pml(), z=td.Boundary.periodic()
    ),
    medium=background,
    shutoff=0,
)

sim.plot(z=0)
plt.show()

# ---------------------------
# Run Simulation Task
# ---------------------------
sim_data = web.run(
    sim,
    task_name=f"2DKerrFC_{freq0*1e-12:.3f}1550_Runtime{run_time*1e12:.3f}_Amp{amp}",
    path="data/simulation_data.hdf5",
    verbose=True,
)

fig, axes = plt.subplots(nrows=4, ncols=1, figsize=(10, 15))
for i in range(4):
    # Extract amplitude data, convert to dB scale, and convert frequency to THz
    monitor = sim_data.monitor_data[f"ring_mode_{i+1}"]
    log_data = 20 * np.log10(
        np.array(monitor.amps.abs.sel(mode_index=1, direction="-"))
    )
    freqs = np.array(monitor.amps.f) / 1e12  # convert Hz to THz

    # Plot vertical lines representing the frequency spectrum
    axes[i].vlines(freqs, ymin=-320, ymax=log_data, color="gray")
    axes[i].set_ylim(-320, -200)
    axes[i].set_title(f"Frequency Spectrum at {sum(pulses[i]) / 2} ps")

    # Remove y-axis tick labels and add y-axis label
    axes[i].set_yticklabels([])
    axes[i].set_ylabel("Power 20 dB / div")

    # Label the x-axis with frequency in THz
    axes[i].set_xlabel("Frequency (THz)")

plt.tight_layout()
plt.show()

import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from IPython.display import display, Image as IPyImage

# ---------------------------
# Measuring round time
# ---------------------------
Round_time = 1.643115e-12  # Duration per segment (s)

# ---------------------------
# Time Data Processing
# ---------------------------
time_data = np.abs(sim_data.monitor_data["out_flux"].flux)
time_data = np.log10(time_data)  # Convert to log scale)

total_duration = sim_data.simulation.run_time  # in seconds
total_samples = len(time_data)
dt = total_duration / total_samples
time_axis = np.linspace(0, total_duration, total_samples)  # in seconds
time_axis_ps = time_axis * 1e12  # convert to ps

# ---------------------------
# Animation Frame Generation with Time Annotation
# ---------------------------
frame_dir = "frames"
os.makedirs(frame_dir, exist_ok=True)

segment_samples = int(total_samples * (Round_time / total_duration))
segment_samples = max(segment_samples, 1)
num_frames = total_samples // segment_samples
print(
    f"Total samples: {total_samples}, Samples per frame: {segment_samples}, Total frames: {num_frames}"
)

frames = []
fig, ax = plt.subplots(figsize=(10, 4), dpi=100)
for i in range(num_frames):
    start_idx = i * segment_samples
    end_idx = start_idx + segment_samples
    if end_idx > total_samples:
        break  # Avoid out-of-bound indices
    segment = time_data[start_idx:end_idx]

    ax.clear()
    ax.plot(segment, color="gray")
    ax.set_ylim(2, 4.5)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_frame_on(False)

    # Calculate current simulation time (in ps) for the frame and quantize to nearest 10 ps
    current_time = (i / (num_frames - 1)) * (run_time * 1e12)
    displayed_time = round(current_time / 10) * 10
    ax.text(
        0.05,
        0.9,
        f"Time: {displayed_time} ps",
        transform=ax.transAxes,
        color="black",
        fontsize=12,
    )

    fig.canvas.draw()
    image = np.array(fig.canvas.renderer.buffer_rgba())
    image_pil = Image.fromarray(image)
    frames.append(image_pil)

plt.close(fig)

gif_path = "time_series_animation.gif"
frames[0].save(gif_path, save_all=True, append_images=frames[1:], duration=100, loop=0)
print(f"GIF generated and saved as {gif_path}")
display(IPyImage(filename=gif_path))

# ---------------------------
# Time-Domain Plot with 4 Normalized Subplots
# ---------------------------
fig, axes = plt.subplots(nrows=4, ncols=1, figsize=(10, 15))
for i, pulse in enumerate(pulses):
    pulse_start_ps = (sum(pulse) - Round_time * 1e12) / 2
    pulse_end_ps = (sum(pulse) + Round_time * 1e12) / 2
    mask = (time_axis_ps >= pulse_start_ps) & (time_axis_ps <= pulse_end_ps)
    segment_time_ps = time_axis_ps[mask]
    segment_data = time_data[mask]

    # Normalize time: 0 at pulse_start and 1 at pulse_end
    normalized_time = (segment_time_ps - pulse_start_ps) / (
        pulse_end_ps - pulse_start_ps
    )

    axes[i].plot(normalized_time, segment_data, color="gray")
    axes[i].set_title(f"Time Domain Signal at {np.mean(pulse):.1f} ps")
    axes[i].set_ylim(2, 4.5)
    axes[i].set_xticks([0, 1])
    axes[i].set_xlabel("Round_time")
    axes[i].set_yticklabels([])
    axes[i].set_ylabel("Flux 10 dB / div")
plt.tight_layout()
plt.show()