from __future__ import annotations

import math
from pathlib import Path

import autograd.numpy as anp
import matplotlib.pyplot as plt
import numpy as np
import tidy3d as td
import tidy3d.web as web
import trimesh
import xarray as xr
from autograd import value_and_grad
from tidy3d.plugins.mode import ModeSolver
from tidy3d.plugins.mode.web import run as run_mode_solver

print(f"Tidy3D version: {td.__version__}")
plt.rcParams.update({"figure.figsize": (9, 5), "axes.grid": True, "grid.alpha": 0.25})

arrow_r_core_um = ARROW_CORE_RADIUS_UM
arrow_r_sio2_1_um = arrow_r_core_um + ARROW_LAYER_SIO2_1_THICKNESS_UM
arrow_r_sin_um = arrow_r_sio2_1_um + ARROW_LAYER_SIN_THICKNESS_UM
arrow_r_sio2_2_um = arrow_r_sin_um + ARROW_LAYER_SIO2_2_THICKNESS_UM

# Concentric layers, largest radius first: later structures override earlier ones where they overlap.
arrow_outer_sio2 = td.Structure(
    geometry=td.Cylinder(center=(0, 0, 0), axis=2, radius=arrow_r_sio2_2_um, length=td.inf),
    medium=sio2_medium,
    name="arrow_outer_sio2",
)
arrow_sin_layer = td.Structure(
    geometry=td.Cylinder(center=(0, 0, 0), axis=2, radius=arrow_r_sin_um, length=td.inf),
    medium=sin_medium,
    name="arrow_sin_layer",
)
arrow_inner_sio2 = td.Structure(
    geometry=td.Cylinder(center=(0, 0, 0), axis=2, radius=arrow_r_sio2_1_um, length=td.inf),
    medium=sio2_medium,
    name="arrow_inner_sio2",
)
arrow_air_core = td.Structure(
    geometry=td.Cylinder(center=(0, 0, 0), axis=2, radius=arrow_r_core_um, length=td.inf),
    medium=background_medium,
    name="arrow_air_core",
)

arrow_mode_extent_um = 2 * arrow_r_sio2_2_um + ARROW_MODE_BUFFER_UM
arrow_mode_sim = td.Simulation(
    center=(0, 0, 0),
    size=(arrow_mode_extent_um, arrow_mode_extent_um, 1.0),
    medium=si_medium,
    structures=[arrow_outer_sio2, arrow_sin_layer, arrow_inner_sio2, arrow_air_core],
    run_time=1e-12,
    grid_spec=td.GridSpec.auto(wavelength=ARROW_WAVELENGTH_UM, min_steps_per_wvl=min_steps_per_wvl),
    boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
)

lambda0_um = ARROW_WAVELENGTH_UM
freq0 = td.C_0 / lambda0_um
freqw = 0.1 * freq0

arrow_mode_plane = td.Box(center=(0, 0, 0), size=(td.inf, td.inf, 0))
arrow_mode_spec = td.ModeSpec(
    num_modes=ARROW_NUM_MODES,
    target_neff=ARROW_TARGET_NEFF,
    precision="double",
    num_pml=ARROW_MODE_NUM_PML,
)
arrow_mode_solver = ModeSolver(
    simulation=arrow_mode_sim,
    plane=arrow_mode_plane,
    mode_spec=arrow_mode_spec,
    freqs=[freq0],
)

if ARROW_MODE_SOLVER_PATH.exists() and not RECOMPUTE_ARROW_MODE:
    arrow_mode_data = td.ModeSolverData.from_file(str(ARROW_MODE_SOLVER_PATH))
    print(f"Loaded cached ARROW mode solver data from {ARROW_MODE_SOLVER_PATH}")
else:
    arrow_mode_data = run_mode_solver(arrow_mode_solver, task_name="arrow_waveguide_mode_solver", verbose=True)
    arrow_mode_data.to_file(str(ARROW_MODE_SOLVER_PATH))
    print(f"Saved ARROW mode solver data to {ARROW_MODE_SOLVER_PATH}")

arrow_n_complex = np.asarray(
    arrow_mode_data.n_complex.sel(f=freq0, method="nearest").values, dtype=complex
).reshape(-1)
ARROW_MODE_INDEX = int(np.argmin(np.abs(arrow_n_complex.real - ARROW_TARGET_NEFF)))
arrow_neff = arrow_n_complex[ARROW_MODE_INDEX]

arrow_e2 = None
arrow_field_template = None
for component_name in E_FIELD_COMPONENTS:
    component = arrow_mode_data.field_components[component_name]
    component = component.sel(mode_index=ARROW_MODE_INDEX, f=freq0, method="nearest").squeeze(drop=True)
    component = component.transpose("x", "y")
    component_e2 = np.abs(component.values) ** 2
    arrow_e2 = component_e2 if arrow_e2 is None else arrow_e2 + component_e2
    if arrow_field_template is None:
        arrow_field_template = component

arrow_field_xy = np.sqrt(arrow_e2)
arrow_field_xy = arrow_field_xy / np.max(arrow_field_xy)
arrow_x_um = np.asarray(arrow_field_template.coords["x"].values, dtype=float)
arrow_y_um = np.asarray(arrow_field_template.coords["y"].values, dtype=float)

arrow = {
    "x_um": arrow_x_um,
    "y_um": arrow_y_um,
    "field_xy": arrow_field_xy,
    "lambda_um": lambda0_um,
    "z_um": 0.0,
}

print(f"lambda0 = {lambda0_um:.4f} um")
print(
    "ARROW layer radii (um): "
    f"core={arrow_r_core_um:.4f}, +SiO2={arrow_r_sio2_1_um:.4f}, "
    f"+SiN={arrow_r_sin_um:.4f}, +SiO2={arrow_r_sio2_2_um:.4f}"
)
print(
    f"Solved {ARROW_NUM_MODES} ARROW modes; selected mode_index={ARROW_MODE_INDEX}, "
    f"n_eff={arrow_neff.real:.6f}, k_eff={arrow_neff.imag:.3e} (target n_eff={ARROW_TARGET_NEFF})"
)
print(f"ARROW x-y target grid = {arrow['field_xy'].shape}")

fig, ax = plt.subplots(figsize=(6, 5))
arrow_mode_solver.plot_field("E", "abs", mode_index=ARROW_MODE_INDEX, f=freq0, ax=ax)
ax.set_title(f"ARROW mode {ARROW_MODE_INDEX}: |E|, n_eff={arrow_neff.real:.4f}")
fig.savefig(RESULTS_DIR / "arrow_profile_from_file.png", dpi=180)
plt.show()

def chord_length_parameterize(points):
    deltas = np.diff(points, axis=0)
    distances = np.sqrt((deltas**2).sum(axis=1))
    cumulative = np.concatenate([[0.0], np.cumsum(distances)])
    return cumulative / cumulative[-1] if cumulative[-1] > 0 else cumulative


def bernstein_matrix(degree, t):
    basis = np.empty((len(t), degree + 1), dtype=float)
    for i in range(degree + 1):
        basis[:, i] = math.comb(degree, i) * (t**i) * ((1.0 - t) ** (degree - i))
    return basis


def fit_bezier_curve(points, degree=3):
    points = np.asarray(points, dtype=float)
    t = chord_length_parameterize(points)
    basis = bernstein_matrix(degree, t)

    control_points = np.zeros((degree + 1, 2), dtype=float)
    control_points[0] = points[0]
    control_points[-1] = points[-1]

    interior_basis = basis[:, 1:-1]
    fixed = np.outer(basis[:, 0], control_points[0]) + np.outer(basis[:, -1], control_points[-1])
    rhs = points - fixed
    control_points[1:-1, 0] = np.linalg.lstsq(interior_basis, rhs[:, 0], rcond=None)[0]
    control_points[1:-1, 1] = np.linalg.lstsq(interior_basis, rhs[:, 1], rcond=None)[0]
    return control_points


def eval_cubic_bezier(cp, u):
    u = anp.reshape(u, (-1,))
    cp = anp.reshape(cp, (4, 2))
    b0 = (1 - u) ** 3
    b1 = 3 * u * (1 - u) ** 2
    b2 = 3 * u**2 * (1 - u)
    b3 = u**3
    return b0[:, None] * cp[0] + b1[:, None] * cp[1] + b2[:, None] * cp[2] + b3[:, None] * cp[3]


def _resample_polyline(points, num_samples):
    points = np.asarray(points, dtype=float)
    deltas = np.diff(points, axis=0)
    arclen = np.concatenate([[0.0], np.cumsum(np.linalg.norm(deltas, axis=1))])
    if arclen[-1] <= 0:
        return np.repeat(points[:1], num_samples, axis=0)
    target = np.linspace(0, arclen[-1], num_samples)
    return np.column_stack([
        np.interp(target, arclen, points[:, 0]),
        np.interp(target, arclen, points[:, 1]),
    ])


def extract_slice_points(vertices_um, y_value, y_tol, num_samples):
    cloud = vertices_um[np.abs(vertices_um[:, 1] - y_value) <= y_tol]
    if len(cloud) < 80:
        return None

    x = cloud[:, 0]
    z = cloud[:, 2]
    bin_count = max(220, num_samples * 3)
    x_bins = np.linspace(x.min(), x.max(), bin_count + 1)
    bin_index = np.digitize(x, x_bins) - 1

    x_centers = []
    z_upper = []
    for idx in range(bin_count):
        in_bin = bin_index == idx
        if np.any(in_bin):
            x_centers.append(float(np.mean(x[in_bin])))
            z_upper.append(float(np.max(z[in_bin])))

    if len(x_centers) < 20:
        return None
    profile = np.column_stack([x_centers, z_upper])
    profile = profile[np.argsort(profile[:, 0])]
    return _resample_polyline(profile, num_samples)


def extract_surface_slices(vertices_um, bounds_um, slice_count, num_samples):
    y_min, y_max = bounds_um[0, 1], bounds_um[1, 1]
    y_span = y_max - y_min
    y_positions = np.linspace(y_min + 0.03 * y_span, y_max - 0.03 * y_span, slice_count)
    base_tol = max((y_positions[1] - y_positions[0]) * 0.45, y_span * 1e-3)

    valid_y = []
    slices = []
    for y_value in y_positions:
        pts = None
        for scale in (1.0, 1.8, 2.8):
            pts = extract_slice_points(vertices_um, y_value, base_tol * scale, num_samples)
            if pts is not None:
                break
        if pts is not None:
            valid_y.append(y_value)
            slices.append(pts)

    if len(slices) != slice_count:
        raise RuntimeError(f"Expected {slice_count} slices, extracted {len(slices)} valid slices.")
    return np.asarray(valid_y), slices


def second_difference_smooth_1d(values, smooth_lambda, fix_first_last=True):
    values = np.asarray(values, dtype=float)
    n = len(values)
    if smooth_lambda <= 0 or n < 4:
        return values.copy()

    d2 = np.zeros((n - 2, n), dtype=float)
    for i in range(n - 2):
        d2[i, i : i + 3] = (1.0, -2.0, 1.0)

    system = np.eye(n) + smooth_lambda * (d2.T @ d2)
    rhs = values.copy()
    if fix_first_last:
        pin_weight = 1e6
        system[0, 0] += pin_weight
        system[-1, -1] += pin_weight
        rhs[0] += pin_weight * values[0]
        rhs[-1] += pin_weight * values[-1]
    return np.linalg.solve(system, rhs)


def smooth_control_points_across_y(control_points_xz):
    smoothed = np.asarray(control_points_xz, dtype=float).copy()
    for point_index in range(4):
        smooth_lambda = SMOOTH_HANDLE_LAMBDA if point_index in (1, 2) else SMOOTH_ENDPOINT_LAMBDA
        for coord_index in range(2):
            smoothed[:, point_index, coord_index] = second_difference_smooth_1d(
                control_points_xz[:, point_index, coord_index],
                smooth_lambda=smooth_lambda,
                fix_first_last=SMOOTH_FIX_FIRST_LAST_SLICE,
            )
    return smoothed


def cross_slice_roughness(control_points_xz):
    d2 = np.diff(np.asarray(control_points_xz), n=2, axis=0)
    return float(np.sqrt(np.mean(d2**2)))


def plot_cross_slice_smoothing(slice_y_um, raw_cp, smooth_cp):
    fig, axes = plt.subplots(2, 2, figsize=(11, 7), constrained_layout=True)
    point_labels = ("P0", "P1", "P2", "P3")
    for ax, point_index in zip(axes.ravel(), range(4)):
        ax.plot(slice_y_um, raw_cp[:, point_index, 1], "o-", color="#adb5bd", linewidth=1.0, markersize=3, label="raw z")
        ax.plot(slice_y_um, smooth_cp[:, point_index, 1], "-", color="#1c7ed6", linewidth=2.0, label="smoothed z")
        ax.set_title(f"{point_labels[point_index]} z(y)")
        ax.set_xlabel("slice y (um)")
        ax.set_ylabel("z (um)")
        ax.grid(True, alpha=0.25)
    axes[0, 0].legend(loc="best")
    fig.savefig(RESULTS_DIR / "bezier_cross_slice_smoothing.png", dpi=180)
    plt.show()


def _subsample_indices(count, max_count, seed=0):
    if count <= max_count:
        return np.arange(count)
    rng = np.random.default_rng(seed)
    return np.sort(rng.choice(count, size=max_count, replace=False))


def _format_bounds(values):
    return f"[{values[0]:.3f}, {values[1]:.3f}]"


def plot_shifted_stl(vertices_um, faces, bounds_um, max_faces=18000, max_points=80000):
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection

    faces = np.asarray(faces, dtype=int)
    face_idx = _subsample_indices(len(faces), max_faces, seed=1)
    point_idx = _subsample_indices(len(vertices_um), max_points, seed=2)
    triangles = vertices_um[faces[face_idx]]
    points = vertices_um[point_idx]

    fig = plt.figure(figsize=(15, 10), constrained_layout=True)
    ax3d = fig.add_subplot(2, 2, 1, projection="3d")
    ax_xz = fig.add_subplot(2, 2, 2)
    ax_xy = fig.add_subplot(2, 2, 3)
    ax_yz = fig.add_subplot(2, 2, 4)

    mesh_artist = Poly3DCollection(
        triangles,
        facecolor="#adb5bd",
        edgecolor="#495057",
        linewidths=0.03,
        alpha=0.80,
    )
    ax3d.add_collection3d(mesh_artist)
    ax3d.set_xlim(bounds_um[0, 0], bounds_um[1, 0])
    ax3d.set_ylim(bounds_um[0, 1], bounds_um[1, 1])
    ax3d.set_zlim(bounds_um[0, 2], bounds_um[1, 2])
    ax3d.set_box_aspect(bounds_um[1] - bounds_um[0])
    ax3d.view_init(elev=24, azim=-55)
    ax3d.set_xlabel("x (um)")
    ax3d.set_ylabel("y (um)")
    ax3d.set_zlabel("z (um)")
    ax3d.set_title(f"shifted STL mesh, z shift = {STL_Z_SHIFT_UM:.2f} um")

    ax_xz.scatter(points[:, 0], points[:, 2], s=0.5, color="#1f77b4", alpha=0.35, rasterized=True)
    ax_xz.set_xlabel("x (um)")
    ax_xz.set_ylabel("z (um)")
    ax_xz.set_title("x-z projection")
    ax_xz.set_aspect("equal", adjustable="box")

    ax_xy.scatter(points[:, 0], points[:, 1], s=0.5, color="#2b8a3e", alpha=0.35, rasterized=True)
    ax_xy.set_xlabel("x (um)")
    ax_xy.set_ylabel("y (um)")
    ax_xy.set_title("x-y footprint")
    ax_xy.set_aspect("equal", adjustable="box")

    ax_yz.scatter(points[:, 1], points[:, 2], s=0.5, color="#f08c00", alpha=0.35, rasterized=True)
    ax_yz.set_xlabel("y (um)")
    ax_yz.set_ylabel("z (um)")
    ax_yz.set_title("y-z projection")
    ax_yz.set_aspect("equal", adjustable="box")

    for ax in (ax_xz, ax_xy, ax_yz):
        ax.grid(True, alpha=0.25)

    fig.savefig(RESULTS_DIR / "stl_shifted_visualization.png", dpi=180)
    plt.show()


mesh_m = trimesh.load_mesh(STL_PATH, force="mesh")
vertices_um = np.asarray(mesh_m.vertices, dtype=float) * 1e6
vertices_um[:, 2] += STL_Z_SHIFT_UM
bounds_um = np.vstack([vertices_um.min(axis=0), vertices_um.max(axis=0)])
print(f"Applied STL z shift: {STL_Z_SHIFT_UM:.3f} um")
print(f"STL vertices: {len(vertices_um)}, faces: {len(mesh_m.faces)}")
print(
    "shifted STL bounds: "
    f"x={_format_bounds(bounds_um[:, 0])} um, "
    f"y={_format_bounds(bounds_um[:, 1])} um, "
    f"z={_format_bounds(bounds_um[:, 2])} um"
)
plot_shifted_stl(vertices_um, mesh_m.faces, bounds_um)

slice_y0_um, slice_points = extract_surface_slices(vertices_um, bounds_um, NUM_SLICES, NUM_PROFILE_POINTS)
control_points0_raw = np.stack([fit_bezier_curve(points, BEZIER_DEGREE) for points in slice_points], axis=0)
control_points0 = smooth_control_points_across_y(control_points0_raw) if SMOOTH_CONTROL_POINTS_ACROSS_Y else control_points0_raw.copy()

y_min0, y_max0 = slice_y0_um.min(), slice_y0_um.max()
raw_roughness = cross_slice_roughness(control_points0_raw)
smoothed_roughness = cross_slice_roughness(control_points0)
max_smoothing_delta = float(np.max(np.abs(control_points0 - control_points0_raw)))

print(f"control_points0 shape: {control_points0.shape}  # slices, 4 points, x/z")
print(f"fixed slice y positions: {slice_y0_um.shape}")
print(f"y span: {y_min0:.3f} to {y_max0:.3f} um")
print(f"cross-slice smoothing enabled: {SMOOTH_CONTROL_POINTS_ACROSS_Y}")
print(f"roughness raw -> smoothed: {raw_roughness:.6f} -> {smoothed_roughness:.6f}")
print(f"max smoothing delta: {max_smoothing_delta:.4f} um")
plot_cross_slice_smoothing(slice_y0_um, control_points0_raw, control_points0)

OPTIMIZED_CONTROL_POINT_INDICES = (1, 2)


def pack_p(control_points_xz):
    control_points_xz = np.asarray(control_points_xz)
    return np.ravel(control_points_xz[:, OPTIMIZED_CONTROL_POINT_INDICES, :])


fixed_control_p0 = anp.array(control_points0[:, 0:1, :])
fixed_control_p3 = anp.array(control_points0[:, 3:4, :])


def unpack_p(p):
    handles = anp.reshape(p, (NUM_SLICES, len(OPTIMIZED_CONTROL_POINT_INDICES), 2))
    return anp.concatenate([fixed_control_p0, handles[:, 0:1, :], handles[:, 1:2, :], fixed_control_p3], axis=1)


p0 = pack_p(control_points0)

cp_min = control_points0.copy()
cp_max = control_points0.copy()
cp_min[:, OPTIMIZED_CONTROL_POINT_INDICES, 0] -= HANDLE_DX_BOUND_UM
cp_max[:, OPTIMIZED_CONTROL_POINT_INDICES, 0] += HANDLE_DX_BOUND_UM
cp_min[:, OPTIMIZED_CONTROL_POINT_INDICES, 1] -= HANDLE_DZ_BOUND_UM
cp_max[:, OPTIMIZED_CONTROL_POINT_INDICES, 1] += HANDLE_DZ_BOUND_UM

p_min = pack_p(cp_min)
p_max = pack_p(cp_max)

print(f"p length: {len(p0)} = 40 * 2 * 2 Bezier handle coordinates")
print(f"optimized points per slice: P1 and P2")
print(f"fixed points per slice: P0 and P3")
print(f"handle x/z bounds: +/-({HANDLE_DX_BOUND_UM}, {HANDLE_DZ_BOUND_UM}) um")
print("slice y positions are fixed and are not part of p")

EXAMPLE_SLICE_INDEX = NUM_SLICES // 2
def plot_single_slice_bounds(
    slice_index=EXAMPLE_SLICE_INDEX,
    control_points_xz=None,
    num_examples=10,
    random_seed=3,
    title_prefix="Initial",
):
    if control_points_xz is None:
        control_points_xz = control_points0

    cp = np.asarray(control_points_xz)[slice_index]
    cp_lo = cp_min[slice_index]
    cp_hi = cp_max[slice_index]
    y_val = slice_y0_um[slice_index]

    rng = np.random.default_rng(random_seed)
    u_plot = np.linspace(0.0, 1.0, 240)

    fig, ax = plt.subplots(figsize=(8.5, 5.6), constrained_layout=True)

    for example_idx in range(num_examples):
        cp_sample = cp.copy()
        for point_index in OPTIMIZED_CONTROL_POINT_INDICES:
            cp_sample[point_index] = cp_lo[point_index] + rng.random(2) * (cp_hi[point_index] - cp_lo[point_index])
        curve_sample = np.asarray(eval_cubic_bezier(cp_sample, u_plot))
        ax.plot(
            curve_sample[:, 0],
            curve_sample[:, 1],
            color="#8da0cb",
            alpha=0.28,
            linewidth=1.2,
            label="example bounded curve" if example_idx == 0 else None,
        )

    curve = np.asarray(eval_cubic_bezier(cp, u_plot))
    ax.plot(curve[:, 0], curve[:, 1], color="black", linewidth=2.6, label=f"{title_prefix.lower()} Bezier curve")
    ax.plot(cp[:, 0], cp[:, 1], "o--", color="#d9480f", linewidth=1.5, markersize=6, label="control polygon")

    point_labels = ["P0", "P1", "P2", "P3"]
    for j, label in enumerate(point_labels):
        if j in OPTIMIZED_CONTROL_POINT_INDICES:
            x0, z0 = cp_lo[j]
            x1, z1 = cp_hi[j]
            ax.add_patch(
                plt.Rectangle(
                    (x0, z0),
                    x1 - x0,
                    z1 - z0,
                    fill=False,
                    edgecolor="#1c7ed6",
                    linewidth=2.0,
                    alpha=0.9,
                )
            )
            suffix = "bounded"
        else:
            suffix = "fixed"
            ax.plot(cp[j, 0], cp[j, 1], marker="s", color="#343a40", markersize=6)
        ax.text(cp[j, 0], cp[j, 1] + 0.12, f"{label} ({suffix})", ha="center", va="bottom", fontsize=9, color="#7a2e0e")

    info = f"fixed slice y = {y_val:.3f} um\np variables: P1/P2 x,z only\nP0/P3 fixed"
    ax.text(
        0.02,
        0.04,
        info,
        transform=ax.transAxes,
        ha="left",
        va="bottom",
        bbox={"boxstyle": "round,pad=0.35", "facecolor": "white", "edgecolor": "#adb5bd", "alpha": 0.92},
    )

    ax.set_title(f"Slice {slice_index}: bounded Bezier handle variation")
    ax.set_xlabel("x (um)")
    ax.set_ylabel("z (um)")
    ax.set_aspect("equal", adjustable="box")
    ax.legend(loc="best")

    print(f"Slice index: {slice_index}")
    print(f"Fixed y: {y_val:.4f} um")
    print("Control-point status for this slice:")
    for j, label in enumerate(point_labels):
        if j in OPTIMIZED_CONTROL_POINT_INDICES:
            print(
                f"  {label}: x [{cp_lo[j, 0]:.4f}, {cp_hi[j, 0]:.4f}] um, "
                f"z [{cp_lo[j, 1]:.4f}, {cp_hi[j, 1]:.4f}] um"
            )
        else:
            print(f"  {label}: fixed at x {cp[j, 0]:.4f} um, z {cp[j, 1]:.4f} um")
    fig.savefig(RESULTS_DIR / "bezier_single_slice_bounds.png", dpi=180)
    plt.show()


plot_single_slice_bounds()

def clip_p_to_bounds(p):
    return anp.minimum(anp.maximum(p, p_min), p_max)


def fixed_y_edges_from_slice_y(slice_y):
    mids = 0.5 * (slice_y[:-1] + slice_y[1:])
    first = slice_y[0] - 0.5 * (slice_y[1] - slice_y[0])
    last = slice_y[-1] + 0.5 * (slice_y[-1] - slice_y[-2])
    return np.concatenate([[first], mids, [last]])


design_y_edges_um = fixed_y_edges_from_slice_y(slice_y0_um)
reflector_y_min_um = design_y_edges_um[0]
reflector_y_max_um = design_y_edges_um[-1]
z_base_um = bounds_um[0, 2] - 0.30
u_poly = np.linspace(0.0, 1.0, NUM_POLYGON_POINTS)

print(f"design slices optimized: {NUM_SLICES}")
print(f"design p length: {len(p0)}")
print(f"design y spacing: {np.diff(design_y_edges_um).mean():.4f} um")


def polygon_from_control_points(cp):
    top = eval_cubic_bezier(cp, u_poly)
    right_base = anp.reshape(anp.stack([top[-1, 0], 0.0 * top[-1, 0] + z_base_um]), (1, 2))
    left_base = anp.reshape(anp.stack([top[0, 0], 0.0 * top[0, 0] + z_base_um]), (1, 2))
    return anp.vstack([top, right_base, left_base])


def make_reflector_structures_from_p(p):
    control_points_xz = unpack_p(p)

    structures = []
    for idx in range(NUM_SLICES):
        vertices_xz = polygon_from_control_points(control_points_xz[idx])
        geom = td.PolySlab(vertices=vertices_xz, axis=1, slab_bounds=(design_y_edges_um[idx], design_y_edges_um[idx + 1]))
        structures.append(
            td.Structure(
                geometry=geom,
                medium=reflector_medium,
                background_medium=background_medium,
                name=f"reflector_slice_{idx:03d}",
            )
        )
    return structures


def make_tip_structure():
    return td.Structure(
        geometry=td.Box(
            center=(tip_x_center_um, tip_y_um, tip_z_um),
            size=(TIP_LENGTH_UM, TIP_SIZE_UM, TIP_SIZE_UM),
        ),
        medium=tip_medium,
        name="horizontal_sin_tip",
    )


reflector_structures0 = make_reflector_structures_from_p(p0)
print(f"Created {len(reflector_structures0)} Bezier PolySlab reflector structures.")

arrow_peak_idx = np.unravel_index(np.argmax(arrow["field_xy"]), arrow["field_xy"].shape)
arrow_peak_file_y_um = float(arrow["y_um"][arrow_peak_idx[1]])

# Horizontal SiN tip placement.
# P0 is the fixed fitted left endpoint for each Bezier slice; use the central slice for z.
tip_y_um = 0.5 * (slice_y0_um.min() + slice_y0_um.max())
left_edge_x_um = float(np.median(control_points0[:, 0, 0]))
left_edge_z_um = float(control_points0[EXAMPLE_SLICE_INDEX, 0, 1])
tip_x_end_um = left_edge_x_um
tip_x_start_um = tip_x_end_um - TIP_LENGTH_UM
tip_x_center_um = 0.5 * (tip_x_start_um + tip_x_end_um)
tip_z_um = WAVEGUIDE_Z_UM
source_x_um = tip_x_end_um - TIP_SOURCE_GAP_FROM_REFLECTOR_UM

# Map the 2D ARROW target into simulation x-y coordinates.
# x is shifted so the aperture geometric center is ARROW_CENTER_X_UM; y peak aligns to the centered tip y.
arrow_raw_x_center_um = 0.5 * (arrow["x_um"].min() + arrow["x_um"].max())
arrow_x_shift_um = ARROW_CENTER_X_UM - arrow_raw_x_center_um
arrow_x_profile_um = arrow["x_um"] + arrow_x_shift_um
arrow_y_profile_um = arrow["y_um"] - arrow_peak_file_y_um + tip_y_um
arrow_xmin, arrow_xmax = arrow_x_profile_um.min(), arrow_x_profile_um.max()
arrow_ymin, arrow_ymax = arrow_y_profile_um.min(), arrow_y_profile_um.max()
reflector_bottom_z_um = z_base_um
arrow_output_z_um = reflector_bottom_z_um - ARROW_MONITOR_GAP_UM
arrow_center = (0.5 * (arrow_xmin + arrow_xmax), 0.5 * (arrow_ymin + arrow_ymax), arrow_output_z_um)
arrow_size = (arrow_xmax - arrow_xmin, arrow_ymax - arrow_ymin, 0.0)

print(f"tip end x = {tip_x_end_um:.4f} um, starts at x = {tip_x_start_um:.4f} um")
print(f"tip center y = {tip_y_um:.4f} um")
print(f"tip / waveguide z = {tip_z_um:.4f} um")
print(f"SiN tip size = {TIP_SIZE_UM:.4f} um square, SiO2 cladding size = {TIP_CLADDING_SIZE_UM:.4f} um square")
print(f"mode source size = {MODE_SOURCE_SIZE_Y_UM:.4f} um by {MODE_SOURCE_SIZE_Z_UM:.4f} um")
print(f"source plane x = {source_x_um:.4f} um, direction = {SOURCE_DIRECTION}x")
print(f"source gap from reflector-side tip end = {TIP_SOURCE_GAP_FROM_REFLECTOR_UM:.4f} um")
print(f"reflector bottom z = {reflector_bottom_z_um:.4f} um")
print(f"ARROW FOM monitor is x-y at z = {arrow_output_z_um:.4f} um")
print(f"ARROW monitor gap below reflector bottom = {ARROW_MONITOR_GAP_UM:.4f} um")
print("No separate ARROW waveguide material box is used; the output region is air background.")
print(f"ARROW aperture center x = {0.5 * (arrow_xmin + arrow_xmax):.4f} um")
print(f"ARROW x shift = {arrow_x_shift_um:.4f} um")
print(f"ARROW FOM x range = [{arrow_xmin:.4f}, {arrow_xmax:.4f}] um")
print(f"ARROW FOM y range = [{arrow_ymin:.4f}, {arrow_ymax:.4f}] um")

source = td.ModeSource(
    center=(source_x_um, tip_y_um, tip_z_um),
    size=(0.0, MODE_SOURCE_SIZE_Y_UM, MODE_SOURCE_SIZE_Z_UM),
    source_time=td.GaussianPulse(freq0=freq0, fwidth=freqw),
    mode_spec=td.ModeSpec(num_modes=1, target_neff=n_tip),
    mode_index=0,
    direction=SOURCE_DIRECTION,
    num_freqs=1,
)

fom_monitor = td.FieldMonitor(
    name=FOM_MONITOR_NAME,
    center=arrow_center,
    size=arrow_size,
    freqs=[freq0],
    fields=E_FIELD_COMPONENTS,
    colocate=True,
)

field_monitor_xy = td.FieldMonitor(
    name="field_xy",
    center=arrow_center,
    size=arrow_size,
    freqs=[freq0],
    fields=E_FIELD_COMPONENTS,
    colocate=True,
)

field_xmin = min(tip_x_start_um, arrow_xmin)
field_xmax = max(tip_x_end_um, arrow_xmax)
field_zmin = min(tip_z_um - MODE_SOURCE_SIZE_Z_UM / 2, arrow_output_z_um)
field_zmax = max(tip_z_um + MODE_SOURCE_SIZE_Z_UM / 2, arrow_output_z_um)
field_monitor_xz = td.FieldMonitor(
    name="field_xz",
    center=(0.5 * (field_xmin + field_xmax), tip_y_um, 0.5 * (field_zmin + field_zmax)),
    size=(field_xmax - field_xmin, 0.0, field_zmax - field_zmin),
    freqs=[freq0],
    fields=E_FIELD_COMPONENTS,
)

opt_xmin = bounds_um[0, 0] - OPT_FIELDS_BUFFER_UM
opt_xmax = bounds_um[1, 0] + OPT_FIELDS_BUFFER_UM
opt_ymin = reflector_y_min_um - OPT_FIELDS_BUFFER_UM
opt_ymax = reflector_y_max_um + OPT_FIELDS_BUFFER_UM
opt_zmin = z_base_um - OPT_FIELDS_BUFFER_UM
opt_zmax = bounds_um[1, 2] + OPT_FIELDS_BUFFER_UM
opt_fields_center = tuple(float(v) for v in (0.5 * (opt_xmin + opt_xmax), 0.5 * (opt_ymin + opt_ymax), 0.5 * (opt_zmin + opt_zmax)))
opt_fields_size = tuple(float(v) for v in (opt_xmax - opt_xmin, opt_ymax - opt_ymin, opt_zmax - opt_zmin))
opt_fields_monitor = td.FieldMonitor(
    name=OPT_FIELDS_MONITOR_NAME,
    center=opt_fields_center,
    size=opt_fields_size,
    freqs=[freq0],
    fields=E_FIELD_COMPONENTS,
    colocate=True,
)

print(f"opt_fields monitor center: {opt_fields_center}")
print(f"opt_fields monitor size:   {opt_fields_size}")


def make_tip_cladding_structure():
    return td.Structure(
        geometry=td.Box(
            center=(tip_x_center_um, tip_y_um, tip_z_um),
            size=(TIP_LENGTH_UM, TIP_CLADDING_SIZE_UM, TIP_CLADDING_SIZE_UM),
        ),
        medium=sio2_medium,
        name="sio2_tip_cladding",
    )


def make_sim(p, include_diagnostics=False):
    reflector_structures = make_reflector_structures_from_p(p)
    tip_cladding_structure = make_tip_cladding_structure()
    tip_structure = make_tip_structure()

    x_min = min(bounds_um[0, 0] - HANDLE_DX_BOUND_UM, arrow_xmin, tip_x_start_um) - pml_buffer_um
    x_max = max(bounds_um[1, 0] + HANDLE_DX_BOUND_UM, arrow_xmax, tip_x_end_um) + pml_buffer_um
    y_min = min(reflector_y_min_um, arrow_ymin, tip_y_um - TIP_CLADDING_SIZE_UM / 2, tip_y_um - MODE_SOURCE_SIZE_Y_UM / 2) - pml_buffer_um
    y_max = max(reflector_y_max_um, arrow_ymax, tip_y_um + TIP_CLADDING_SIZE_UM / 2, tip_y_um + MODE_SOURCE_SIZE_Y_UM / 2) + pml_buffer_um
    z_min = min(z_base_um, arrow_output_z_um, tip_z_um - TIP_CLADDING_SIZE_UM / 2, tip_z_um - MODE_SOURCE_SIZE_Z_UM / 2) - pml_buffer_um
    z_max = max(bounds_um[1, 2] + HANDLE_DZ_BOUND_UM, arrow_output_z_um, tip_z_um + TIP_CLADDING_SIZE_UM / 2, tip_z_um + MODE_SOURCE_SIZE_Z_UM / 2) + pml_buffer_um

    sim_center = ((x_min + x_max) / 2, (y_min + y_max) / 2, (z_min + z_max) / 2)
    sim_size = (x_max - x_min, y_max - y_min, z_max - z_min)

    monitors = [fom_monitor]
    if include_diagnostics:
        monitors.extend([field_monitor_xy, field_monitor_xz, opt_fields_monitor])

    return td.Simulation(
        center=sim_center,
        size=sim_size,
        medium=background_medium,
        structures=[tip_cladding_structure, tip_structure] + reflector_structures,
        sources=[source],
        monitors=monitors,
        run_time=run_time_factor / freqw,
        grid_spec=td.GridSpec.auto(wavelength=lambda0_um, min_steps_per_wvl=min_steps_per_wvl),
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
    )


sim0 = make_sim(p0, include_diagnostics=True)
print(f"Simulation size: {sim0.size}")
print(f"Structures: {len(sim0.structures)}")
print(f"Diagnostic monitors: {[mnt.name for mnt in sim0.monitors]}")

# Plot the geometry and the simulation setup separately using Tidy3D's own plotting.


def plot_tidy3d_cuts(source_alpha, monitor_alpha, title_prefix, file_name):
    fig, axes = plt.subplots(1, 3, figsize=(17, 4.8), constrained_layout=True)

    sim0.plot_eps(
        y=tip_y_um,
        freq=freq0,
        ax=axes[0],
        source_alpha=source_alpha,
        monitor_alpha=monitor_alpha,
    )
    axes[0].set_title(f"{title_prefix}: x-z cut at y = {tip_y_um:.2f} um")

    sim0.plot_eps(
        z=arrow_output_z_um,
        freq=freq0,
        ax=axes[1],
        source_alpha=source_alpha,
        monitor_alpha=monitor_alpha,
    )
    axes[1].set_title(f"{title_prefix}: x-y cut at ARROW plane")

    sim0.plot_eps(
        x=source_x_um,
        freq=freq0,
        ax=axes[2],
        source_alpha=source_alpha,
        monitor_alpha=monitor_alpha,
    )
    axes[2].set_title(f"{title_prefix}: y-z cut at source x")

    fig.savefig(RESULTS_DIR / file_name, dpi=180)
    plt.show()


plot_tidy3d_cuts(
    source_alpha=0.0,
    monitor_alpha=0.0,
    title_prefix="Tidy3D structures only",
    file_name="tidy3d_structures_only.png",
)

plot_tidy3d_cuts(
    source_alpha=0.7,
    monitor_alpha=0.7,
    title_prefix="Tidy3D simulation setup",
    file_name="tidy3d_simulation_setup.png",
)

arrow_target_da = xr.DataArray(
    arrow["field_xy"],
    coords={"x": arrow_x_profile_um, "y": arrow_y_profile_um},
    dims=("x", "y"),
)


def _target_on_monitor_grid(field_on_plane):
    x_grid = np.asarray(field_on_plane.coords["x"].values, dtype=float)
    y_grid = np.asarray(field_on_plane.coords["y"].values, dtype=float)
    target_interp = arrow_target_da.interp(x=x_grid, y=y_grid).fillna(0.0)
    return np.asarray(target_interp.values, dtype=float)


def electric_field_magnitude_on_monitor(sim_data, monitor_name, dims_order=None):
    e2 = None
    template = None
    monitor_data = sim_data[monitor_name]

    for component_name in E_FIELD_COMPONENTS:
        component = monitor_data.field_components[component_name]
        component = component.sel(f=freq0, method="nearest").squeeze(drop=True)
        if dims_order is not None:
            component = component.transpose(*dims_order)
        elif {"x", "y"}.issubset(component.dims):
            component = component.transpose("x", "y")

        values = component.values
        component_e2 = anp.abs(values) ** 2
        e2 = component_e2 if e2 is None else e2 + component_e2
        if template is None:
            template = component

    return anp.sqrt(e2), template


def electric_field_magnitude_dataarray(sim_data, monitor_name, dims_order=None):
    magnitude, template = electric_field_magnitude_on_monitor(sim_data, monitor_name, dims_order=dims_order)
    coords = {dim: template.coords[dim] for dim in template.dims}
    return xr.DataArray(np.asarray(magnitude), coords=coords, dims=template.dims)


def measure_arrow_profile_match(sim_data):
    simulated, field_on_plane = electric_field_magnitude_on_monitor(
        sim_data, FOM_MONITOR_NAME, dims_order=("x", "y")
    )
    target = anp.array(_target_on_monitor_grid(field_on_plane))

    eps = 1e-18
    target = target / anp.sqrt(anp.sum(anp.abs(target) ** 2) + eps)
    simulated = simulated / anp.sqrt(anp.sum(anp.abs(simulated) ** 2) + eps)
    overlap = anp.sum(anp.conj(target) * simulated)
    return anp.real(anp.conj(overlap) * overlap)


def objective(p, step_num=0, verbose=False):
    p = clip_p_to_bounds(p)
    sim = make_sim(p, include_diagnostics=False)
    sim_data = web.run(
        sim,
        task_name=f"bezier_reflector_arrow_step_{step_num:03d}",
        path=str(RESULTS_DIR / f"bezier_reflector_arrow_step_{step_num:03d}.hdf5"),
        verbose=verbose,
    )
    return measure_arrow_profile_match(sim_data)


objective_and_grad = value_and_grad(objective)
print(f"Scalar FOM uses x-y monitor '{FOM_MONITOR_NAME}' and total electric-field magnitude |E| from {E_FIELD_COMPONENTS}.")

RUN_COST_ESTIMATE = True

if RUN_COST_ESTIMATE:
    cost_sim = make_sim(p0, include_diagnostics=False)
    task_id = web.upload(cost_sim, task_name="bezier_reflector_arrow_cost_check", verbose=False)
    estimated_cost = web.estimate_cost(task_id)
    print(f"Estimated cost: {estimated_cost:.4f} FlexCredits")
else:
    print("Set RUN_COST_ESTIMATE = True after geometry plots are approved.")

RUN_OPTIMIZATION = True

num_steps = 40
learning_rate = 0.3

# Lightweight visualization uses optimization data already computed each step.
VISUALIZE_OPTIMIZATION = True
VISUALIZE_OPTIMIZATION_EVERY = 1

# Optional field diagnostics require extra cloud simulations. Keep False unless the cost is intentional.
RUN_DIAGNOSTIC_FIELDS_DURING_OPTIMIZATION = False
DIAGNOSTIC_FIELDS_EVERY = 5


def _as_numpy_control_points(p_current):
    return np.asarray(unpack_p(np.asarray(p_current, dtype=float)), dtype=float)


def _handle_array_from_p(p_current):
    return np.asarray(p_current, dtype=float).reshape(NUM_SLICES, len(OPTIMIZED_CONTROL_POINT_INDICES), 2)


def plot_handle_movement_progress(step, p_current):
    handle_delta = _handle_array_from_p(p_current - p0)

    fig, ax = plt.subplots(figsize=(7, 5), constrained_layout=True)
    labels = ("P1", "P2")
    for handle_idx, label in enumerate(labels):
        ax.plot(slice_y0_um, handle_delta[:, handle_idx, 0], "-", linewidth=1.5, label=f"{label}.x")
        ax.plot(slice_y0_um, handle_delta[:, handle_idx, 1], "--", linewidth=1.5, label=f"{label}.z")
    ax.axhline(0.0, color="black", linewidth=0.8)
    ax.set_xlabel("slice y (um)")
    ax.set_ylabel("offset from initial (um)")
    ax.set_title(f"optimized handle movement, step {step:03d}")
    ax.legend(loc="best", fontsize=8, ncol=2)
    ax.grid(True, alpha=0.25)
    fig.savefig(RESULTS_DIR / f"handle_movement_step_{step:03d}.png", dpi=180)
    fig.savefig(RESULTS_DIR / "handle_movement_latest.png", dpi=180)
    plt.show()


def plot_reflector_optimization_progress(step, p_current, J_history, grad_norm_history, update_norm_history):
    cp_current = _as_numpy_control_points(p_current)
    handle_delta = _handle_array_from_p(p_current - p0)
    bound_half_width = _handle_array_from_p((p_max - p_min) / 2)
    normalized_delta = np.divide(
        handle_delta,
        bound_half_width,
        out=np.zeros_like(handle_delta),
        where=bound_half_width > 0,
    )
    max_bound_usage = float(np.max(np.abs(normalized_delta)))

    fig, axes = plt.subplots(2, 2, figsize=(13, 9), constrained_layout=True)
    axes = axes.ravel()
    u_plot = np.linspace(0.0, 1.0, 180)
    slice_indices = np.unique(np.linspace(0, NUM_SLICES - 1, min(7, NUM_SLICES), dtype=int))
    colors = plt.cm.viridis(np.linspace(0.0, 1.0, len(slice_indices)))

    for color, slice_index in zip(colors, slice_indices):
        curve0 = np.asarray(eval_cubic_bezier(control_points0[slice_index], u_plot))
        curve_current = np.asarray(eval_cubic_bezier(cp_current[slice_index], u_plot))
        axes[0].plot(curve0[:, 0], curve0[:, 1], color="#adb5bd", linewidth=1.0, alpha=0.75)
        axes[0].plot(
            curve_current[:, 0],
            curve_current[:, 1],
            color=color,
            linewidth=1.8,
            label=f"y={slice_y0_um[slice_index]:.2f} um",
        )
    axes[0].set_aspect("equal", adjustable="box")
    axes[0].set_xlabel("x (um)")
    axes[0].set_ylabel("z (um)")
    axes[0].set_title("Bezier reflector slices: gray initial, color current")
    axes[0].legend(loc="best", fontsize=7)

    if J_history:
        axes[1].plot(np.arange(1, len(J_history) + 1), J_history, "o-", color="#1c7ed6")
    axes[1].set_xlabel("evaluated step")
    axes[1].set_ylabel("FOM")
    axes[1].set_title("scalar ARROW |E| profile match")
    axes[1].grid(True, alpha=0.25)

    labels = ("P1", "P2")
    for handle_idx, label in enumerate(labels):
        axes[2].plot(slice_y0_um, handle_delta[:, handle_idx, 0], "-", linewidth=1.5, label=f"{label}.x")
        axes[2].plot(slice_y0_um, handle_delta[:, handle_idx, 1], "--", linewidth=1.5, label=f"{label}.z")
    axes[2].axhline(0.0, color="black", linewidth=0.8)
    axes[2].set_xlabel("slice y (um)")
    axes[2].set_ylabel("offset from initial (um)")
    axes[2].set_title("optimized handle movement")
    axes[2].legend(loc="best", fontsize=8, ncol=2)
    axes[2].grid(True, alpha=0.25)

    if grad_norm_history:
        axes[3].semilogy(np.arange(1, len(grad_norm_history) + 1), grad_norm_history, "o-", label="grad norm")
    if update_norm_history:
        axes[3].semilogy(np.arange(1, len(update_norm_history) + 1), update_norm_history, "s-", label="update norm")
    axes[3].axhline(max(max_bound_usage, 1e-16), color="#f08c00", linestyle=":", label="max bound usage")
    axes[3].set_xlabel("step")
    axes[3].set_title(f"optimizer scales, max bound usage = {max_bound_usage:.3f}")
    axes[3].legend(loc="best", fontsize=8)
    axes[3].grid(True, alpha=0.25)

    latest_j = J_history[-1] if J_history else np.nan
    fig.suptitle(f"Bezier reflector inverse-design progress: step {step:03d}, FOM={latest_j:.6e}")
    fig.savefig(RESULTS_DIR / f"optimization_step_{step:03d}.png", dpi=180)
    fig.savefig(RESULTS_DIR / "optimization_latest.png", dpi=180)
    plt.show()


def plot_iteration_field_diagnostics(step, sim_data):
    field_xy_e = electric_field_magnitude_dataarray(sim_data, "field_xy", dims_order=("x", "y"))
    field_xz_e = electric_field_magnitude_dataarray(sim_data, "field_xz", dims_order=("x", "z"))
    fom_e = electric_field_magnitude_dataarray(sim_data, FOM_MONITOR_NAME, dims_order=("x", "y"))

    fig, axes = plt.subplots(1, 3, figsize=(16, 4), constrained_layout=True)
    field_xy_e.T.plot(ax=axes[0], cmap="magma")
    axes[0].set_title(f"step {step:03d}: |E| x-y diagnostic")
    field_xz_e.T.plot(ax=axes[1], cmap="magma")
    axes[1].set_title(f"step {step:03d}: |E| x-z slice")
    fom_e.T.plot(ax=axes[2], cmap="magma")
    axes[2].set_title(f"step {step:03d}: |E| ARROW plane")
    fig.savefig(RESULTS_DIR / f"optimization_fields_step_{step:03d}.png", dpi=180)
    plt.show()


p = p0.copy()
J_history = []
grad_norm_history = []
update_norm_history = []
p_history = [p.copy()]

if RUN_OPTIMIZATION:
    from tidy3d.plugins.autograd import adam, apply_updates

    optimizer = adam(learning_rate=learning_rate)
    opt_state = optimizer.init(p)

    for step in range(1, num_steps + 1):
        p_eval = np.array(p)
        value, grad = objective_and_grad(p_eval, step_num=step, verbose=False)
        grad = np.asarray(grad, dtype=float)
        grad_norm = float(np.linalg.norm(grad))

        updates, opt_state = optimizer.update(-grad, opt_state, p_eval)  # negative gradient because we maximize J
        p_next = apply_updates(p_eval, updates)
        p_next = clip_p_to_bounds(p_next)
        update_norm = float(np.linalg.norm(np.asarray(p_next) - p_eval))

        J_history.append(float(value))
        grad_norm_history.append(grad_norm)
        update_norm_history.append(update_norm)
        p_history.append(np.array(p_next))

        print(
            f"step={step:03d} J={float(value):.6e} "
            f"grad_norm={grad_norm:.6e} update_norm={update_norm:.6e}"
        )

        if VISUALIZE_OPTIMIZATION and step % VISUALIZE_OPTIMIZATION_EVERY == 0:
            plot_handle_movement_progress(step, np.asarray(p_next))

        if RUN_DIAGNOSTIC_FIELDS_DURING_OPTIMIZATION and step % DIAGNOSTIC_FIELDS_EVERY == 0:
            sim_diag = make_sim(np.asarray(p_next), include_diagnostics=True)
            sim_data_diag = web.run(
                sim_diag,
                task_name=f"bezier_reflector_arrow_diagnostic_step_{step:03d}",
                path=str(RESULTS_DIR / f"bezier_reflector_arrow_diagnostic_step_{step:03d}.hdf5"),
                verbose=False,
            )
            plot_iteration_field_diagnostics(step, sim_data_diag)

        p = np.asarray(p_next)

    if VISUALIZE_OPTIMIZATION:
        plot_reflector_optimization_progress(step, np.asarray(p), J_history, grad_norm_history, update_norm_history)
else:
    print("Set RUN_OPTIMIZATION = True only after cost approval.")