import autograd.numpy as anp
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import time
import tidy3d as td
import tidy3d.web as web
from tidy3d.exceptions import WebError
from tidy3d.plugins.autograd import adam, optimize

td.config.logging.level = "ERROR"

# Materials.
air_index = 1.0
gan_index = 2.4

air_medium = td.Medium(permittivity=air_index**2)
gan_medium = td.Medium(permittivity=gan_index**2)

# Center wavelength and source bandwidth.
center_wavelength = 0.5
bandwidth_wavelength = 0.1

# Physical lateral size of the finite GaN micro LED pixel.
pixel_size_xy = 3.5
pixel_half_width = pixel_size_xy / 2.0

# Add one center wavelength of air padding on each side before the PML.
domain_size_xy = pixel_size_xy + 2.0 * center_wavelength

# Etched GaN layer. The air-hole pattern occupies this vertical interval.
gan_z_max = 0.5
etch_z_min = 0.3
etch_z_max = 0.5
etch_depth = etch_z_max - etch_z_min
etch_center_z = 0.5 * (etch_z_min + etch_z_max)

# Vertical source, flux monitor, and boundary placement.
bottom_source_spacing = 0.5
structure_monitor_spacing = 0.4
monitor_top_spacing = 0.25

field_monitor_z = gan_z_max + structure_monitor_spacing
source_z = -field_monitor_z
domain_z_min = source_z - bottom_source_spacing
domain_z_max = field_monitor_z + monitor_top_spacing
domain_size_z = domain_z_max - domain_z_min
domain_center_z = 0.5 * (domain_z_min + domain_z_max)
domain_size = (domain_size_xy, domain_size_xy, domain_size_z)
domain_center = (0.0, 0.0, domain_center_z)
gan_z_min = domain_z_min
# Extend the lower GaN substrate past the bottom PML so it does not terminate there.
gan_substrate_z_min = domain_z_min - 2.0 * center_wavelength
source_center = (0.0, 0.0, source_z)
source_norm_box_size = 0.4

# Single center frequency sampled by each flux monitor.
wavelength_min = center_wavelength - bandwidth_wavelength / 2.0
wavelength_max = center_wavelength + bandwidth_wavelength / 2.0

freq0 = td.C_0 / center_wavelength
freq_min = td.C_0 / wavelength_max
freq_max = td.C_0 / wavelength_min
freq_width = freq_max - freq_min
monitor_freqs = np.array([freq0])

# Time window and base mesh density used by the optimization simulations.
run_time = 200.0 / freq_width
min_steps_per_wvl = 15

def make_grid_spec(steps_per_wvl=min_steps_per_wvl):
    """Create an automatic grid with enforced refinement around the etched layer."""
    dl = wavelength_min / gan_index / steps_per_wvl
    grid_spec = td.GridSpec.auto(
        wavelength=center_wavelength,
        min_steps_per_wvl=steps_per_wvl,
        override_structures=[td.MeshOverrideStructure(
            geometry=td.Box.from_bounds(
                rmin=(-pixel_half_width - 0.5 * wavelength_min, -pixel_half_width - 0.5 * wavelength_min, etch_z_min - 0.5 * wavelength_min),
                rmax=(pixel_half_width + 0.5 * wavelength_min, pixel_half_width + 0.5 * wavelength_min, etch_z_max + 0.5 * wavelength_min),
            ),
            dl=(dl, dl, dl),
            enforce=True,
        )]
    )
    return grid_spec


optimization_steps_per_wvl = min_steps_per_wvl
analysis_steps_per_wvl = 2 * optimization_steps_per_wvl

grid_spec = make_grid_spec(optimization_steps_per_wvl)
analysis_grid_spec = make_grid_spec(analysis_steps_per_wvl)

param_names = ("pitch_a", "pitch_b", "r_a", "r_b", "theta_deg")

pitch_min = 0.25
pitch_max = 0.60
radius_min = 0.05
radius_fraction_max = 0.40
theta_max_deg = 60.0
fixed_dx = 0.0
fixed_dy = 0.0
operation = "intersection"

initial_physical_param_values = {
    "pitch_a": 0.45,
    "pitch_b": 0.50,
    "r_a": 0.10,
    "r_b": 0.15,
    "theta_deg": 21.8,
}
initial_physical_params = np.array(
    [initial_physical_param_values[name] for name in param_names],
    dtype=float,
)

sqrt3 = np.sqrt(3.0)

def build_fixed_lattice_index_window():
    """Return fixed candidate hexagonal lattice indices covering the pixel."""
    cover_radius = np.sqrt(2.0) * pixel_half_width + 2.0 * pitch_max
    j_limit = int(np.ceil(cover_radius / ((sqrt3 / 2.0) * pitch_min))) + 2
    i_limit = int(np.ceil(cover_radius / pitch_min)) + 2
    radial_index_limit = np.sqrt(2.0) * pixel_half_width / pitch_min + radius_fraction_max

    return tuple(
        (i_idx, j_idx)
        for j_idx in range(-j_limit, j_limit + 1)
        for i_idx in range(-i_limit, i_limit + 1)
        if np.sqrt(i_idx**2 + i_idx * j_idx + j_idx**2) <= radial_index_limit
    )


fixed_lattice_index_pairs = build_fixed_lattice_index_window()

def pack_params_named(params):
    """Map normalized design variables to physical geometry parameters."""
    pitch_a = pitch_min + params[0] * (pitch_max - pitch_min)
    pitch_b = pitch_min + params[1] * (pitch_max - pitch_min)

    radius_a_max = radius_fraction_max * pitch_a
    radius_b_max = radius_fraction_max * pitch_b

    radius_a = radius_min + params[2] * (radius_a_max - radius_min)
    radius_b = radius_min + params[3] * (radius_b_max - radius_min)

    theta_deg = params[4] * theta_max_deg

    return {
        "pitch_a": pitch_a,
        "pitch_b": pitch_b,
        "r_a": radius_a,
        "r_b": radius_b,
        "theta_deg": theta_deg,
        "dx": fixed_dx,
        "dy": fixed_dy,
    }


def encode_initial_params(params):
    """Convert physical starting parameters to normalized optimizer variables."""
    pitch_a, pitch_b, radius_a, radius_b, theta_deg = np.asarray(params, dtype=float)
    normalized = np.empty(len(param_names), dtype=float)

    normalized[0] = (pitch_a - pitch_min) / (pitch_max - pitch_min)
    normalized[1] = (pitch_b - pitch_min) / (pitch_max - pitch_min)

    radius_a_max = radius_fraction_max * pitch_a
    radius_b_max = radius_fraction_max * pitch_b
    normalized[2] = (radius_a - radius_min) / (radius_a_max - radius_min)
    normalized[3] = (radius_b - radius_min) / (radius_b_max - radius_min)

    normalized[4] = theta_deg / theta_max_deg

    return np.clip(normalized, 0.0, 1.0)


def project_params(params):
    """Clip optimizer variables to the normalized design interval."""
    return np.clip(np.asarray(params, dtype=float).copy(), 0.0, 1.0)


def format_params(params):
    """Format normalized variables as physical parameter values."""
    values = pack_params_named(np.asarray(params, dtype=float))
    return ", ".join(f"{name}={value:.4f}" for name, value in values.items())

def lattice_center(pitch, i_idx, j_idx, theta_rad=0.0, dx=0.0, dy=0.0):
    """Compute one rotated and shifted hexagonal lattice-site center."""
    x_base = pitch * (i_idx + 0.5 * j_idx)
    y_base = pitch * (sqrt3 / 2.0) * j_idx

    cos_theta = anp.cos(theta_rad)
    sin_theta = anp.sin(theta_rad)
    x_rot = cos_theta * x_base - sin_theta * y_base
    y_rot = sin_theta * x_base + cos_theta * y_base
    return x_rot + dx, y_rot + dy


def build_hex_lattice_holes(pitch, radius, theta_rad=0.0, dx=0.0, dy=0.0):
    """Build one hexagonal lattice of cylindrical air holes."""
    cylinders = []
    for i_idx, j_idx in fixed_lattice_index_pairs:
        x_pos, y_pos = lattice_center(
            pitch,
            i_idx,
            j_idx,
            theta_rad=theta_rad,
            dx=dx,
            dy=dy,
        )
        cylinders.append(
            td.Cylinder(
                center=(x_pos, y_pos, etch_center_z),
                radius=radius,
                length=etch_depth,
                axis=2,
            )
        )

    return td.GeometryGroup(geometries=tuple(cylinders))


def combine_lattices(lattice_a, lattice_b, operation):
    """Combine two hole lattices with the selected boolean operation."""
    if operation == "intersection":
        return lattice_a & lattice_b
    if operation == "symmetric_difference":
        return lattice_a ^ lattice_b
    raise ValueError("operation must be 'intersection' or 'symmetric_difference'.")


def build_etch_geometry(params, operation=operation):
    """Create the clipped air-etch geometry from normalized design variables."""
    named = pack_params_named(params)
    theta_rad = named["theta_deg"] * anp.pi / 180.0

    lattice_a = build_hex_lattice_holes(named["pitch_a"], named["r_a"])
    lattice_b = build_hex_lattice_holes(
        named["pitch_b"],
        named["r_b"],
        theta_rad=theta_rad,
        dx=named["dx"],
        dy=named["dy"],
    )

    raw_pattern = combine_lattices(lattice_a, lattice_b, operation=operation)
    pixel_mask = td.Box.from_bounds(
        rmin=(-pixel_half_width, -pixel_half_width, etch_z_min),
        rmax=(pixel_half_width, pixel_half_width, etch_z_max),
    )

    return raw_pattern & pixel_mask

def make_source(source_center=source_center, polarization="Ex"):
    """Create a broadband point dipole source."""
    return td.PointDipole(
        center=source_center,
        polarization=polarization,
        source_time=td.GaussianPulse(
            freq0=freq0,
            fwidth=freq_width,
            amplitude=1.0,
        ),
    )


def make_top_flux_monitor():
    """Create the upward flux monitor with adjoint flux support enabled."""
    return td.FluxMonitor(
        center=(0.0, 0.0, field_monitor_z),
        size=(domain_size[0], domain_size[1], 0.0),
        freqs=tuple(monitor_freqs.tolist()),
        name="top_flux_monitor",
        normal_dir="+",
        enable_adjoint=True,
    )


def make_norm_flux_monitor(source_center=source_center):
    """Create the small adjoint-capable flux box for emitted dipole power."""
    return td.FluxMonitor(
        center=source_center,
        size=(source_norm_box_size, source_norm_box_size, source_norm_box_size),
        freqs=tuple(monitor_freqs.tolist()),
        name="norm_flux_monitor",
        enable_adjoint=True,
    )

def build_static_gan_structures():
    """Create the lower GaN block and upper GaN cap."""
    lower_gan = td.Structure(
        geometry=td.Box.from_bounds(
            rmin=(-pixel_half_width, -pixel_half_width, gan_substrate_z_min),
            rmax=(pixel_half_width, pixel_half_width, etch_z_min),
        ),
        medium=gan_medium,
        name="gan_lower",
    )

    upper_gan_cap = td.Structure(
        geometry=td.Box.from_bounds(
            rmin=(-pixel_half_width, -pixel_half_width, etch_z_min),
            rmax=(pixel_half_width, pixel_half_width, etch_z_max),
        ),
        medium=gan_medium,
        name="gan_cap",
    )

    return lower_gan, upper_gan_cap


def build_traced_etch_structure(etch_geometry):
    """Wrap a raw or processed air etch geometry as a structure inside GaN."""
    return td.Structure(
        geometry=etch_geometry,
        medium=air_medium,
        background_medium=gan_medium,
        name="etch_holes",
    )


def build_simulation(
    params,
    monitors,
    source_center,
    source_polarization,
    grid_spec,
    etch_geometry=None,
    include_etch=True,
):
    """Assemble a finite-pixel micro LED simulation."""
    structures = list(build_static_gan_structures())
    if include_etch:
        if etch_geometry is None:
            etch_geometry = build_etch_geometry(params, operation=operation)
        structures.append(build_traced_etch_structure(etch_geometry))

    return td.Simulation(
        center=domain_center,
        size=domain_size,
        medium=air_medium,
        structures=tuple(structures),
        sources=(make_source(source_center=source_center, polarization=source_polarization),),
        monitors=monitors,
        boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=True),
        grid_spec=grid_spec,
        run_time=run_time,
    )


def build_optimization_simulation(params):
    """Build the differentiable simulation used inside the adjoint objective."""
    return build_simulation(
        params,
        monitors=(
            make_top_flux_monitor(),
            make_norm_flux_monitor(),
        ),
        source_center=source_center,
        source_polarization="Ex",
        grid_spec=grid_spec,
    )


def build_analysis_simulation(
    params,
    grid_spec,
    source_center=source_center,
    source_polarization="Ex",
    etch_geometry=None,
    include_etch=True,
):
    """Build a standard FDTD simulation for normalized top LEE evaluation."""
    monitors = (
        make_top_flux_monitor(),
        make_norm_flux_monitor(source_center=source_center),
    )

    return build_simulation(
        params,
        monitors=monitors,
        source_center=source_center,
        source_polarization=source_polarization,
        grid_spec=grid_spec,
        etch_geometry=etch_geometry,
        include_etch=include_etch,
    )


def build_interface_reference_simulation(
    grid_spec,
    source_center=source_center,
    source_polarization="Ex",
):
    """Build the flat top finite-pixel reference simulation."""
    return td.Simulation(
        center=domain_center,
        size=domain_size,
        medium=air_medium,
        structures=(
            td.Structure(
                geometry=td.Box.from_bounds(
                    rmin=(-pixel_half_width, -pixel_half_width, gan_substrate_z_min),
                    rmax=(pixel_half_width, pixel_half_width, gan_z_max),
                ),
                medium=gan_medium,
                name="gan_interface_reference",
            ),
        ),
        sources=(make_source(source_center=source_center, polarization=source_polarization),),
        monitors=(
            make_top_flux_monitor(),
            make_norm_flux_monitor(source_center=source_center),
        ),
        boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=True),
        grid_spec=grid_spec,
        run_time=run_time,
    )

params0 = project_params(encode_initial_params(initial_physical_params))
sim0 = build_optimization_simulation(params0)
sim0.validate_pre_upload()

print(format_params(params0))

fig, ax = plt.subplots(1, 2, figsize=(11, 4), tight_layout=True)

sim0.plot_eps(z=etch_center_z, ax=ax[0])
ax[0].set_xlim(-pixel_half_width, pixel_half_width)
ax[0].set_ylim(-pixel_half_width, pixel_half_width)
ax[0].set_aspect("equal")
ax[0].set_title("initial etch pattern")

sim0.plot_eps(y=0, ax=ax[1])
ax[1].set_xlim(-pixel_half_width, pixel_half_width)
ax[1].set_ylim(gan_z_min, field_monitor_z)
ax[1].set_title("vertical cross section")

plt.show()

optimization_folder_name = "ClipOperationLightExtractorOptimization"
objective_max_retries = 3
objective_retry_sleep = 30


def objective(params, task_name="clip-operation-light-extractor"):
    """Run one adjoint solve and return normalized top LEE."""
    sim = build_optimization_simulation(params)
    for attempt in range(objective_max_retries + 1):
        try:
            sim_data = web.run(
                sim,
                task_name=task_name,
                folder_name=optimization_folder_name,
                verbose=False,
            )
            break
        except WebError:
            if attempt == objective_max_retries:
                raise
            time.sleep(objective_retry_sleep)

    top_flux = sim_data["top_flux_monitor"].flux.values
    emitted_flux = sim_data["norm_flux_monitor"].flux.values
    return anp.squeeze(top_flux / emitted_flux)

num_steps = 30
learning_rate = 0.01
print_every = 5

optimizer = adam(learning_rate=learning_rate)
optimization_step = {"index": 0}

params_history = []
physical_params_history = []


def objective_for_current_step(current_params):
    step_index = optimization_step["index"]
    task_name = f"clip-operation-light-extractor-iter-{step_index:02d}"
    return objective(current_params, task_name=task_name)


def optimization_callback(current_params, gradient, opt_state, step_index, value):
    current_params = project_params(current_params)
    gradient_norm = np.linalg.norm(np.asarray(gradient, dtype=float))

    params_history.append(current_params.copy())
    physical_params_history.append(pack_params_named(current_params))

    if step_index == 0 or (step_index + 1) % print_every == 0 or step_index == num_steps - 1:
        print(f"step = {step_index + 1:02d}")
        print(f"  top LEE = {100 * value:.2f}%")
        print(f"  grad_norm = {gradient_norm:.4e}")
        print(f"  {format_params(current_params)}")

    optimization_step["index"] = step_index + 1

params_final, opt_state, optimization_history = optimize(
    objective_for_current_step,
    np.array(params0),
    optimizer,
    num_steps,
    bounds=(0.0, 1.0),
    callback=optimization_callback,
    direction="max",
)
params_final = project_params(params_final)

objective_history = optimization_history["objective_fn_val"]
gradient_norm_history = optimization_history["grad_norm"]
params_history.append(params_final.copy())
physical_params_history.append(pack_params_named(params_final))

fig, ax = plt.subplots(1, 2, figsize=(11, 4), tight_layout=True)

ax[0].plot(np.arange(1, num_steps + 1), 100 * np.asarray(objective_history), "o-")
ax[0].set_xlabel("iteration")
ax[0].set_ylabel("top LEE (%)")
ax[0].set_title("normalized objective")

ax[1].semilogy(np.arange(1, num_steps + 1), gradient_norm_history, "o-")
ax[1].set_xlabel("iteration")
ax[1].set_ylabel("gradient norm")
ax[1].set_title("gradient norm")

plt.show()

sim_final = build_optimization_simulation(params_final)

fig, ax = plt.subplots(1, 2, figsize=(11, 5), tight_layout=True)

sim0.plot_eps(z=etch_center_z, ax=ax[0])
ax[0].set_xlim(-pixel_half_width, pixel_half_width)
ax[0].set_ylim(-pixel_half_width, pixel_half_width)
ax[0].set_aspect("equal")
ax[0].set_title("initial")

sim_final.plot_eps(z=etch_center_z, ax=ax[1])
ax[1].set_xlim(-pixel_half_width, pixel_half_width)
ax[1].set_ylim(-pixel_half_width, pixel_half_width)
ax[1].set_aspect("equal")
ax[1].set_title("optimized")

plt.show()

pd.DataFrame(
    [
        {"design": "initial", **pack_params_named(params0)},
        {"design": "optimized", **pack_params_named(params_final)},
    ]
).round(4)

import shapely


minimum_fabrication_feature = 0.05
use_cleaned_design_for_analysis = True
pixel_mask_polygon = shapely.box(
    -pixel_half_width,
    -pixel_half_width,
    pixel_half_width,
    pixel_half_width,
)


def replace_infinite_coordinates(coords):
    """Replace infinite Shapely coordinates by a finite box outside the pixel."""
    coords = np.asarray(coords, dtype=float)
    finite_bound = 10.0 * domain_size_xy
    return np.where(np.isinf(coords), np.sign(coords) * finite_bound, coords)


def finite_shapely_geometry(shapely_geometry):
    """Return a Shapely geometry with infinite coordinates replaced."""
    bounds = np.asarray(shapely_geometry.bounds, dtype=float)
    if shapely_geometry.is_empty or not np.any(np.isinf(bounds)):
        return shapely_geometry
    return shapely.transform(
        shapely_geometry,
        replace_infinite_coordinates,
        include_z=False,
    )


def polygon_list(shapely_geometry):
    """Flatten a Shapely geometry into nonempty polygons."""
    polygons = []
    valid_geometry = shapely.make_valid(finite_shapely_geometry(shapely_geometry))
    for part in shapely.get_parts(valid_geometry):
        if part.is_empty:
            continue
        if part.geom_type == "Polygon" and part.area > 0.0:
            polygons.append(part)
        elif part.geom_type in {"MultiPolygon", "GeometryCollection"}:
            polygons.extend(polygon_list(part))
    return polygons


def as_multipolygon(shapely_geometry):
    """Normalize a Shapely object into a valid `MultiPolygon`."""
    polygons = polygon_list(shapely_geometry)
    return shapely.MultiPolygon(polygons)


def constrain_to_pixel(shapely_geometry):
    """Clip etch regions to the finite LED pixel footprint."""
    return as_multipolygon(shapely_geometry.intersection(pixel_mask_polygon))


def etch_geometry_to_multipolygon(etch_geometry, cleanup=False, quad_segs=24):
    """Slice a Tidy3D etch geometry into clipped 2D Shapely polygons."""
    polygons = etch_geometry.intersections_plane(
        z=etch_center_z,
        cleanup=cleanup,
        quad_segs=quad_segs,
    )
    if len(polygons) == 0:
        return shapely.MultiPolygon([])

    section = shapely.union_all([finite_shapely_geometry(poly) for poly in polygons])
    return constrain_to_pixel(section)


def filter_small_components(multipolygon, minimum_feature):
    """Remove etched components below an equivalent diameter threshold."""
    minimum_area = np.pi * (0.5 * minimum_feature) ** 2
    polygons = [
        polygon
        for polygon in polygon_list(multipolygon)
        if polygon.area >= minimum_area
    ]
    return shapely.MultiPolygon(polygons)


def clean_etch_multipolygon(multipolygon, minimum_feature):
    """Apply opening followed by closing to a 2D etch pattern."""
    radius = 0.5 * minimum_feature
    cleaned = constrain_to_pixel(as_multipolygon(multipolygon))
    cleaned = constrain_to_pixel(cleaned.buffer(-radius, quad_segs=8))
    cleaned = constrain_to_pixel(cleaned.buffer(radius, quad_segs=8))
    cleaned = constrain_to_pixel(cleaned.buffer(radius, quad_segs=8))
    cleaned = constrain_to_pixel(cleaned.buffer(-radius, quad_segs=8))
    cleaned = filter_small_components(as_multipolygon(cleaned), minimum_feature)
    return constrain_to_pixel(cleaned)


def ring_vertices(ring):
    """Convert a Shapely ring to vertices without the duplicated endpoint."""
    x_coords, y_coords = ring.coords.xy
    return list(zip(np.asarray(x_coords[:-1]), np.asarray(y_coords[:-1])))


def polygon_to_polyslab(polygon):
    """Convert one processed Shapely polygon into Tidy3D geometry."""
    exterior = td.PolySlab(
        vertices=ring_vertices(polygon.exterior),
        slab_bounds=(etch_z_min, etch_z_max),
        axis=2,
    )

    if len(polygon.interiors) == 0:
        return exterior

    interiors = tuple(
        td.PolySlab(
            vertices=ring_vertices(interior),
            slab_bounds=(etch_z_min, etch_z_max),
            axis=2,
        )
        for interior in polygon.interiors
        if len(interior.coords) >= 4
    )

    if not interiors:
        return exterior

    return exterior - td.GeometryGroup(geometries=interiors)


def multipolygon_to_etch_geometry(multipolygon):
    """Convert processed Shapely etch regions back to Tidy3D geometry."""
    polygons = [
        polygon
        for polygon in polygon_list(multipolygon)
        if not polygon.is_empty and polygon.area > 0.0
    ]
    if len(polygons) == 0:
        return None

    polyslabs = tuple(polygon_to_polyslab(polygon) for polygon in polygons)
    if len(polyslabs) == 1:
        return polyslabs[0]
    return td.GeometryGroup(geometries=polyslabs)


def summarize_multipolygon(multipolygon, label):
    """Summarize component count and equivalent feature sizes."""
    polygons = polygon_list(multipolygon)
    areas = np.asarray([polygon.area for polygon in polygons], dtype=float)
    equivalent_diameters = 2.0 * np.sqrt(areas / np.pi) if len(areas) else np.array([])

    return {
        "geometry": label,
        "num_components": len(polygons),
        "total_area_um2": float(np.sum(areas)) if len(areas) else 0.0,
        "min_equiv_diameter_um": float(np.min(equivalent_diameters)) if len(areas) else np.nan,
        "median_equiv_diameter_um": float(np.median(equivalent_diameters)) if len(areas) else np.nan,
    }


raw_etch_multipolygons = {
    "initial": etch_geometry_to_multipolygon(
        build_etch_geometry(params0, operation=operation),
        cleanup=False,
    ),
    "optimized": etch_geometry_to_multipolygon(
        build_etch_geometry(params_final, operation=operation),
        cleanup=False,
    ),
}
cleaned_etch_multipolygons = {
    design: clean_etch_multipolygon(raw_etch, minimum_fabrication_feature)
    for design, raw_etch in raw_etch_multipolygons.items()
}
cleaned_etch_geometries = {
    design: multipolygon_to_etch_geometry(cleaned_etch)
    for design, cleaned_etch in cleaned_etch_multipolygons.items()
}

raw_initial_etch = raw_etch_multipolygons["initial"]
raw_optimized_etch = raw_etch_multipolygons["optimized"]
cleaned_initial_etch_multipolygon = cleaned_etch_multipolygons["initial"]
cleaned_optimized_etch_multipolygon = cleaned_etch_multipolygons["optimized"]
cleaned_initial_etch_geometry = cleaned_etch_geometries["initial"]
cleaned_optimized_etch_geometry = cleaned_etch_geometries["optimized"]

cleanup_summary_records = []
for design, raw_etch in raw_etch_multipolygons.items():
    for label, etch in (
        ("raw", raw_etch),
        ("processed", cleaned_etch_multipolygons[design]),
    ):
        record = summarize_multipolygon(etch, label)
        record["design"] = design
        cleanup_summary_records.append(record)

cleanup_summary = pd.DataFrame(cleanup_summary_records)[
    [
        "design",
        "geometry",
        "num_components",
        "total_area_um2",
        "min_equiv_diameter_um",
        "median_equiv_diameter_um",
    ]
]
cleanup_summary.round(4)

def plot_shapely_geometry(ax, geometry, facecolor, edgecolor="black", alpha=0.8, label=None):
    """Draw Shapely polygons on a Matplotlib axis."""
    for polygon in polygon_list(geometry):
        x_coords, y_coords = polygon.exterior.xy
        ax.fill(x_coords, y_coords, facecolor=facecolor, edgecolor=edgecolor, alpha=alpha, label=label)
        label = None
        for interior in polygon.interiors:
            x_hole, y_hole = interior.xy
            ax.fill(x_hole, y_hole, facecolor="white", edgecolor=edgecolor, alpha=1.0)


def plot_cleanup_comparison(axes, design, raw_etch, processed_etch):
    """Plot raw, processed, and changed etch regions for one design."""
    removed_etch = as_multipolygon(raw_etch.difference(processed_etch))
    added_etch = as_multipolygon(processed_etch.difference(raw_etch))

    plot_shapely_geometry(axes[0], raw_etch, facecolor="#4e79a7", alpha=0.85)
    axes[0].set_title(f"{design}: raw etch")

    plot_shapely_geometry(axes[1], processed_etch, facecolor="#59a14f", alpha=0.85)
    axes[1].set_title(f"{design}: processed etch")

    plot_shapely_geometry(axes[2], raw_etch, facecolor="#bab0ac", alpha=0.35)
    plot_shapely_geometry(axes[2], removed_etch, facecolor="#e15759", edgecolor="#e15759", alpha=0.9, label="removed")
    plot_shapely_geometry(axes[2], added_etch, facecolor="#4e79a7", edgecolor="#4e79a7", alpha=0.9, label="filled")
    axes[2].set_title(f"{design}: cleanup change")
    axes[2].legend(loc="upper right")


fig, ax = plt.subplots(2, 3, figsize=(13, 7), tight_layout=True)

plot_cleanup_comparison(
    ax[0],
    "initial",
    raw_initial_etch,
    cleaned_initial_etch_multipolygon,
)
plot_cleanup_comparison(
    ax[1],
    "optimized",
    raw_optimized_etch,
    cleaned_optimized_etch_multipolygon,
)

for axis in ax.ravel():
    axis.set_xlim(-pixel_half_width, pixel_half_width)
    axis.set_ylim(-pixel_half_width, pixel_half_width)
    axis.set_aspect("equal")
    axis.set_xlabel("x (um)")
    axis.set_ylabel("y (um)")

plt.show()

initial_processed_sim = build_analysis_simulation(
    params0,
    grid_spec=analysis_grid_spec,
    etch_geometry=cleaned_initial_etch_geometry,
    include_etch=cleaned_initial_etch_geometry is not None,
)
optimized_processed_sim = build_analysis_simulation(
    params_final,
    grid_spec=analysis_grid_spec,
    etch_geometry=cleaned_optimized_etch_geometry,
    include_etch=cleaned_optimized_etch_geometry is not None,
)
initial_processed_sim.validate_pre_upload()
optimized_processed_sim.validate_pre_upload()

initial_raw_sim = build_analysis_simulation(params0, grid_spec=analysis_grid_spec)
optimized_raw_sim = build_analysis_simulation(params_final, grid_spec=analysis_grid_spec)

fig, ax = plt.subplots(2, 2, figsize=(11, 9), tight_layout=True)

initial_raw_sim.plot_eps(z=etch_center_z, ax=ax[0, 0])
ax[0, 0].set_title("initial raw")

initial_processed_sim.plot_eps(z=etch_center_z, ax=ax[0, 1])
ax[0, 1].set_title("initial processed")

optimized_raw_sim.plot_eps(z=etch_center_z, ax=ax[1, 0])
ax[1, 0].set_title("optimized raw")

optimized_processed_sim.plot_eps(z=etch_center_z, ax=ax[1, 1])
ax[1, 1].set_title("optimized processed")

for axis in ax.ravel():
    axis.set_xlim(-pixel_half_width, pixel_half_width)
    axis.set_ylim(-pixel_half_width, pixel_half_width)
    axis.set_aspect("equal")

plt.show()

def summarize_lee_result(sim_data, design, polarization, position_index=0, source_center=source_center):
    """Package one LEE result into a dataframe row."""
    emitted_flux = np.asarray(sim_data["norm_flux_monitor"].flux.values, dtype=float)
    extracted_flux = np.asarray(sim_data["top_flux_monitor"].flux.values, dtype=float)
    top_lee = float(np.squeeze(extracted_flux / emitted_flux))

    return {
        "design": design,
        "polarization": polarization,
        "position_index": position_index,
        "x": source_center[0],
        "y": source_center[1],
        "z": source_center[2],
        "top_lee": top_lee,
    }