# Import relevant libraries
import numpy as np
import tidy3d as td
import tidy3d.web as web
from shapely.geometry import Polygon as ShapelyPolygon
import gdstk
import matplotlib.pyplot as plt
import scipy.integrate as integrate
from scipy.optimize import fsolve
from scipy.optimize import brentq

# Define simulation parameters
wl0 = 0.638
freq0 = td.C_0 / wl0
fwidth = 0.2 * freq0
npoints = 100
freqs = np.linspace(freq0 - fwidth / 2, freq0 + fwidth / 2, npoints)
run_time = 1e-12

# Define material properties
n_SiN = 2.038
n_SiO2 = 1.4569
SiN = td.Medium(permittivity=n_SiN**2)
SiO2 = td.Medium(permittivity=n_SiO2**2)

# Define geometric parameters
h_SiN = 0.15
W_wg = 0.5
R = 15 # effective bend radius
theta_t = np.pi / 2 # total bend angle
L_str = 3.0 # straight input/output section length
inf = 100
A = 2.4
buffer = 1

# Poly3 Functions
def poly3_theta(s, Rc, tt, tc):
    """Tangential angle along Poly3 (Eq. 1)."""
    return (2 * Rc * (tt - tc) * s**3 - s**4) / (2 * Rc**4 * (tt - tc)**3)


def _poly3_endpoint(Rc, tt, tc, n=3000):
    """Integrate (cos, sin) to find the Poly3 endpoint."""
    s = np.linspace(0, Rc * (tt - tc), n)
    th = poly3_theta(s, Rc, tt, tc)
    return np.trapezoid(np.cos(th), s), np.trapezoid(np.sin(th), s)


def _constraint(Rc, R, tt, tc):
    """Bisector symmetry constraint (Eq. 2)."""
    px, py = _poly3_endpoint(Rc, tt, tc)
    ang = (tt - tc) / 2
    xo = px - Rc * np.sin(ang)
    yo = py + Rc * np.cos(ang)
    return xo * np.cos(tt / 2) + (yo - R) * np.sin(tt / 2)


def find_Rc(R, tt, tc):
    """Solve for the Circ arc radius Rc given (R, theta_t, theta_c)."""
    return brentq(_constraint, 1e-4 * R, 10 * R, args=(R, tt, tc))


def generate_baseline(R, tt, tc, n=1200):
    """
    Build the Poly3 + Circ + Poly3' centreline.
    Starts at (0, 0) heading in +x; ends at (R, R) heading in +y (for theta_t = pi/2).
    """
    Rc  = find_Rc(R, tt, tc)
    L_p = Rc * (tt - tc)

    # Poly3
    n_p  = n // 3
    s_p  = np.linspace(0, L_p, n_p)
    th_p = poly3_theta(s_p, Rc, tt, tc)
    xp   = np.concatenate([[0.0], np.cumsum(np.cos(th_p[:-1]) * np.diff(s_p))])
    yp   = np.concatenate([[0.0], np.cumsum(np.sin(th_p[:-1]) * np.diff(s_p))])

    # Circ arc
    ang     = (tt - tc) / 2
    xo      = xp[-1] - Rc * np.sin(ang)
    yo      = yp[-1] + Rc * np.cos(ang)
    n_c     = n // 3
    a_start = np.arctan2(yp[-1] - yo, xp[-1] - xo)
    alphas  = np.linspace(a_start, a_start + tc, n_c)
    xc      = xo + Rc * np.cos(alphas)
    yc      = yo + Rc * np.sin(alphas)

    # Poly3' (mirror of Poly3)
    n_pp  = n - n_p - n_c + 2
    s_pp  = np.linspace(0, L_p, n_pp)
    th_pp = tt - poly3_theta(L_p - s_pp, Rc, tt, tc)
    xpp   = np.concatenate([[xc[-1]], xc[-1] + np.cumsum(np.cos(th_pp[:-1]) * np.diff(s_pp))])
    ypp   = np.concatenate([[yc[-1]], yc[-1] + np.cumsum(np.sin(th_pp[:-1]) * np.diff(s_pp))])

    return (np.concatenate([xp, xc[1:], xpp[1:]]),
            np.concatenate([yp, yc[1:], ypp[1:]]))


def offset_curve(x, y, d):
    """
    Offset a curve by signed distance d along its left-pointing normal.
    d > 0 -> inward (toward centre of curvature)
    d < 0 -> outward
    """
    dx, dy = np.gradient(x), np.gradient(y)
    mag    = np.hypot(dx, dy) + 1e-30
    return x + d * (-dy / mag), y + d * (dx / mag)

# Euler Functions
def euler_curve(L_max, A, num_points=50):
    Ls = np.linspace(0, L_max, num_points)

    y1 = np.array([integrate.quad(lambda theta: A * np.sin(theta**2 / 2), 0, L / A)[0] for L in Ls])
    x1 = np.array([integrate.quad(lambda theta: A * np.cos(theta**2 / 2), 0, L / A)[0] for L in Ls])

    return x1, y1

def calculate_L_max(R_eff, A, num_points=50, tolerance=1e-12):
    def func(L_max):
        L_max = float(L_max.item())
        x1, y1 = euler_curve(L_max, A, num_points)

        # compute the derivative at L_max
        k = -(x1[-1] - x1[-2]) / (y1[-1] - y1[-2])
        xp = x1[-1]
        yp = y1[-1]

        # check if the derivative is continuous at L_max
        R = np.sqrt(
            ((R_eff + k * xp - yp) / (k + 1) - xp) ** 2
            + (-(R_eff + k * xp - yp) / (k + 1) + R_eff - yp) ** 2
        )
        return np.abs(R - A**2 / L_max)

    res = fsolve(func, x0=(1,), full_output=True)

    error = res[1]["fvec"]

    if error > tolerance:
        raise ValueError(f"No solutions for given A: {A}.")
    L_max = res[0][0]

    return L_max

L_max = calculate_L_max(R, A)

x1, y1 = euler_curve(L_max, A)

xp = x1[-1]
yp = y1[-1]

R_min = A**2 / L_max

# Circular Functions
def circle(var):
    a = var[0]
    b = var[1]
    Func = np.empty(2)
    Func[0] = (xp - a) ** 2 + (yp - b) ** 2 - R_min**2
    Func[1] = (R - yp - a) ** 2 + (R - xp - b) ** 2 - R_min**2
    return Func

def mirror_path(x_points, y_points, mirror_y=0):
    y_mirrored = 2 * mirror_y - y_points
    return x_points, y_mirrored

def line_to_structure(x, y, w, t):
    cell = gdstk.Cell("bend")  # define a gds cell

    # add points to include the input and output straight waveguides

    cell.add(gdstk.FlexPath(x + 1j * y, w, layer=1, datatype=0))  # add path to cell

    # define structure from cell
    bend = td.Structure(
        geometry=td.PolySlab.from_gds(
            cell,
            gds_layer=1,
            axis=2,
            slab_bounds=(-t/2, t/2),
        )[0],
        medium=SiN,
    )

    return bend

def poly_slab(verts, med):
    return td.Structure(
        geometry=td.PolySlab(vertices=verts, slab_bounds=(-h_SiN/2, h_SiN/2), axis=2, sidewall_angle=np.radians(0)),
        medium=med,
    )

# Visualization Functions
def plot_conversion_efficiency(sim_data, mon_name=["mode"], labels=[], title="", show=True):
    # extract the conversion efficiency from the mode monitor
    for mon in mon_name:
        amp = sim_data[mon].amps.sel(mode_index=0, direction="-")
        freq = sim_data[mon].amps.f
        wl = td.C_0 / freq
        T = np.abs(amp) ** 2
        plt.plot(wl, -10*np.log10(np.abs(T)**2), label=labels[mon_name.index(mon)])

    plt.xlim(0.613, 0.663)
    plt.xlabel(r"Wavelength ($\mu m$)")
    plt.ylabel("Bending Loss (dB)")
    plt.legend()
    plt.title(title)
    if show:
        plt.show()

def plot_field(sim_data):

    ax = sim_data.plot_field(
        field_monitor_name="field", field_name="E", val="abs", f= freq0, vmin=0, vmax=100,
    )
    plt.show()