import numpy as np
import matplotlib.pyplot as plt

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

# ---------------------------
# Parameter Definitions
# ---------------------------
# Geometrical parameters
wg_width = 1.25
couple_width = 0.5
ring_radius = 10
ring_wg_width = wg_width
wg_spacing = 2.5
buffer = 5.0

# Simulation dimensions
x_span = 2 * wg_spacing + 2 * ring_radius + 2 * buffer
y_span = 2 * ring_radius + ring_wg_width + wg_width + couple_width + 2 * buffer
sim_center_y = (couple_width + wg_width) / 2
wg_insert_x = ring_radius + wg_spacing
wg_center_y = ring_radius + ring_wg_width / 2.0 + couple_width + wg_width / 2.0

# Wavelength and frequency parameters
freq0 = 195.835e12
amp = 18
run_time = 400e-12
min_steps_per_wvl = 30
fwidth = 44.9689e12

# ---------------------------
# Medium Definitions
# ---------------------------
n_bg = 1.0
n_solid = 1.5
background = td.Medium(permittivity=n_bg**2)
solid = td.Medium(permittivity=n_solid**2)

n_kerr_2 = 2.0e-8
kerr_chi3 = 4 * (n_solid**2) * td.constants.EPSILON_0 * td.constants.C_0 * n_kerr_2 / 3
chi3_model = td.NonlinearSpec(
    models=[td.NonlinearSusceptibility(chi3=kerr_chi3)], num_iters=10
)
kerr_solid = td.Medium(permittivity=n_solid**2, nonlinear_spec=chi3_model)

# ---------------------------
# Structure and Simulation Region Definitions
# ---------------------------
# Define structure
waveguide = td.Structure(
    geometry=td.Box(
        center=[0, wg_center_y, 0],
        size=[td.inf, wg_width, td.inf],
    ),
    medium=solid,
    name="waveguide",
)
# Ring Structure
outer = td.Cylinder(
    center=[0, 0, 0],
    axis=2,
    radius=ring_radius + ring_wg_width / 2.0,
    length=td.inf,
)
inner = td.Cylinder(
    center=[0, 0, 0],
    axis=2,
    radius=ring_radius - ring_wg_width / 2.0,
    length=td.inf,
)
ring = td.Structure(
    geometry=outer - inner,
    medium=kerr_solid,
    name="outer_ring",
)
# Define mode plane and ring mode plane
mode_plane = td.Box(center=[-wg_insert_x, wg_center_y, 0], size=[0, 10, td.inf])
ring_mode_plane = td.Box(center=[0, -ring_radius, 0], size=[0, 10, td.inf])

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import argrelextrema


# ----------------------------
# Small utilities (internal)
# ----------------------------
def _to_numpy(x):
    """Return a NumPy array from NumPy/xarray-like input."""
    return np.asarray(x.values) if hasattr(x, "values") else np.asarray(x)


def _get_time_axis(time_data, total_time=None):
    """
    Extract time coordinate from an xarray-like object if available.
    Otherwise, build a uniform time axis using `total_time`.
    """
    # Prefer xarray coords: time_data.coords['t'] or ['time']
    if hasattr(time_data, "coords"):
        for key in ("t", "time"):
            if key in time_data.coords:
                return np.asarray(time_data.coords[key].values)

    # Fallback: time_data.t (some objects expose this)
    if hasattr(time_data, "t"):
        try:
            return np.asarray(time_data.t.values)
        except Exception:
            pass

    # Last resort: uniform axis
    if total_time is None:
        raise ValueError(
            "No time coordinate found in `time_data`. Provide `total_time` (seconds) "
            "or pass an object containing time coords (e.g., time_data.coords['t'])."
        )
    y = _to_numpy(time_data)
    return np.linspace(0.0, total_time, len(y), endpoint=False)


# ----------------------------
# 1) Round-trip integrated flux
# ----------------------------
def plot_roundtrip_flux(
    flux_data, run_time=400, roundtrip_time=0.3310882353, line_color="b"
):
    """
    Integrate (sum) the flux within each round-trip window and plot versus time.

    Notes
    -----
    - Integration uses a simple summation over samples (same as original).
    - Time axis is built as uniform sampling over [0, run_time) (same as original).
    """
    y = _to_numpy(flux_data)

    n = len(y)
    t = np.linspace(0, run_time, n, endpoint=False)

    n_rt = int(run_time / roundtrip_time)

    rt_times = []
    rt_flux = []

    for i in range(n_rt):
        t0 = i * roundtrip_time
        t1 = (i + 1) * roundtrip_time
        tc = 0.5 * (t0 + t1)

        i0 = np.searchsorted(t, t0)
        i1 = np.searchsorted(t, t1)

        if i1 > i0:
            rt_times.append(tc)
            rt_flux.append(np.sum(y[i0:i1]))

    rt_times = np.asarray(rt_times)
    rt_flux = np.asarray(rt_flux)

    fig, ax = plt.subplots(figsize=(9, 3.3), dpi=100)
    ax.plot(rt_times, rt_flux, color=line_color, linewidth=1.5)

    ax.set_yticks([])
    ax.set_xlabel("Time (ps)", fontsize=12)
    ax.set_ylabel("Intracavity Energy (a.u.)", fontsize=12)

    plt.tight_layout()
    plt.show()


# ----------------------------
# 2) Frequency spectrum plot
# ----------------------------
def plot_frequency_spectrum(
    monitor_data,
    time_ps=395,
    mode_index=1,
    direction="-",
    ylim=(-320, -200),
    figsize=(9, 3.3),
    color="gray",
):
    """
    Plot a line spectrum (vlines) from monitor data at a given time label.

    Assumes
    -------
    monitor_data.amps.abs.sel(mode_index=..., direction=...) is available,
    and monitor_data.amps.f provides frequencies in Hz.
    """
    fig, ax = plt.subplots(figsize=figsize)

    amps = np.array(
        monitor_data.amps.abs.sel(mode_index=mode_index, direction=direction)
    )
    log_data = 20 * np.log10(amps)
    freqs_thz = np.array(monitor_data.amps.f) / 1e12

    ax.vlines(freqs_thz, ymin=ylim[0], ymax=log_data, color=color)
    ax.set_ylim(ylim)
    ax.set_title(f"Frequency Spectrum at {time_ps} ps")

    ax.set_yticklabels([])
    ax.set_ylabel("Power 20 dB / div")
    ax.set_xlabel("Frequency (THz)")

    plt.tight_layout()
    plt.show()


# ----------------------------
# 3) Local maxima in a time window
# ----------------------------
def plot_time_window_maxima(
    time_data,
    start_time,
    end_time,
    roundtrip_time=0.3310882353,
    min_height=None,
    min_distance=3,
    total_time=None,
    time_unit_for_plot="ps",
    ylim=(0, 1e5),
):
    """
    Plot raw time data in a window and overlay a curve through local maxima.

    Parameters
    ----------
    time_data : array-like or xarray-like
        Time series samples. If it contains time coordinates ('t' or 'time'), they are used.
        Otherwise provide `total_time` (seconds) to build a uniform axis.
    start_time, end_time : float
        Window bounds in seconds (same as original code behavior).
    min_height : float or None
        Optional threshold for maxima values.
    min_distance : int
        Neighborhood size (in points) for argrelextrema(order=...).
    time_unit_for_plot : {"s","ns","ps","fs"}
        Display unit for the time axis.
    ylim : tuple or None
        y-axis limits.
    """
    if end_time <= start_time:
        raise ValueError("end_time must be greater than start_time.")

    y = _to_numpy(time_data)
    t = _get_time_axis(time_data, total_time=total_time)

    i0 = int(np.searchsorted(t, start_time, side="left"))
    i1 = int(np.searchsorted(t, end_time, side="left"))
    i0 = max(0, min(i0, len(y)))
    i1 = max(0, min(i1, len(y)))

    y_win = y[i0:i1]
    t_win = t[i0:i1]

    max_idx = argrelextrema(y_win, np.greater, order=min_distance)[0]
    if min_height is not None and len(max_idx) > 0:
        max_idx = max_idx[y_win[max_idx] >= min_height]

    t_max = t_win[max_idx]
    y_max = y_win[max_idx]

    unit_scale = {"s": 1.0, "ns": 1e-9, "ps": 1e-12, "fs": 1e-15}
    if time_unit_for_plot not in unit_scale:
        raise ValueError("time_unit_for_plot must be one of: 's', 'ns', 'ps', 'fs'.")

    scale = unit_scale[time_unit_for_plot]
    t_win_disp = t_win / scale
    t_max_disp = t_max / scale

    start_disp = start_time / scale
    end_disp = end_time / scale
    _ = roundtrip_time / scale  # kept for parity; not used in plotting (same effect)

    plt.figure(figsize=(9, 3.3))
    plt.plot(
        t_win_disp,
        y_win,
        color="lightgray",
        alpha=1,
        linewidth=0.5,
        label="flux Data",
    )

    if len(t_max) > 1:
        # Keep the exact original behavior: spline only if > 5 maxima, else connect points.
        if len(t_max) > 5:
            from scipy.interpolate import make_interp_spline

            xnew = np.linspace(t_max_disp[0], t_max_disp[-1], 300)
            spl = make_interp_spline(t_max_disp, y_max, k=3)
            ynew = spl(xnew)
            plt.plot(xnew, ynew, alpha=0.7, linewidth=2, label="pulse envelope")
        else:
            plt.plot(t_max_disp, y_max, alpha=0.7, linewidth=2, label="pulse envelope")

    plt.title(
        f"Time window: {start_disp:.6g}{time_unit_for_plot} to {end_disp:.6g}{time_unit_for_plot}"
    )
    plt.xlabel(f"Time ({time_unit_for_plot})")
    plt.ylabel("Flux (a.u.)")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.yticks([])

    if ylim is not None:
        plt.ylim(*ylim)

    plt.tight_layout()
    plt.show()