# Substrate
substrate = td.Structure(
    geometry=td.Box.from_bounds(
        rmin=(-td.inf, -td.inf, -1000),
        rmax=(+td.inf, +td.inf, z_center_slab - thickness / 2),
    ),
    medium=td.Medium(permittivity=permittivity),
)

# Source
src = td.PlaneWave(
    center=(0, 0, -Lz / 2 + buffer),
    size=(td.inf, td.inf, 0),
    source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
    direction="+",
)

# Monitor used in the cost function
diffractionmonitor = td.DiffractionMonitor(
    name="diffractionmonitor",
    center=(0, 0, +Lz / 2 - buffer),
    size=(td.inf, td.inf, 0),
    interval_space=[1, 1, 1],
    colocate=False,
    freqs=[freq0],
    apodization=td.ApodizationSpec(),
    normal_dir="+",
)

# Creating filters
radius = 0.1

filter_project = make_filter_and_project(radius, dl_design_region)
erosion_dilation_penalty = make_erosion_dilation_penalty(radius, dl_design_region, beta=10)


def get_eps(params: anp.ndarray, beta: float) -> anp.ndarray:
    """Get the permittivity values (1, permittivity) array as a function of the parameters (0, 1)"""
    density = filter_project(params, beta)
    eps = rescale(density, 1, permittivity)
    return eps.reshape((nx, ny, 1))


def make_slab(params: anp.ndarray, beta: float) -> td.Structure:
    """Make the optimization design region"""
    box = td.Box(center=(0, 0, z_center_slab), size=(2 * length, 2 * length, thickness))
    eps_data = get_eps(params, beta)
    return td.Structure.from_permittivity_array(geometry=box, eps_data=eps_data)

def make_sim(params: anp.ndarray, beta: float) -> td.Simulation:
    """The simulation as a function of the design parameters."""

    slab = make_slab(params, beta)

    # Mesh override structure to ensure uniform dl across the slab
    design_region_mesh = td.MeshOverrideStructure(
        geometry=slab.geometry,
        dl=[dl_design_region] * 3,
        enforce=False,
    )

    return td.Simulation(
        size=(length, length, Lz),
        grid_spec=td.GridSpec.auto(
            min_steps_per_wvl=min_steps_per_wvl,
            override_structures=[design_region_mesh],
        ),
        boundary_spec=td.BoundarySpec(
            x=td.Boundary(
                plus=td.Periodic(),
                minus=td.Periodic(),
            ),
            y=td.Boundary(
                plus=td.Periodic(),
                minus=td.Periodic(),
            ),
            z=td.Boundary(
                plus=td.PML(),
                minus=td.PML(),
            ),
        ),
        structures=[substrate, slab],
        monitors=[diffractionmonitor],
        sources=[src],
        run_time=run_time,
    )

# Make symmetric, random starting parameters
params0 = np.random.random((nx, ny))
params0 += np.fliplr(params0)
params0 += np.flipud(params0)
params0 /= 4.0

beta0 = 1.0

sim = make_sim(params=params0, beta=beta0)

f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), tight_layout=True)
ax1 = sim.plot_eps(x=0, ax=ax1)
ax2 = sim.plot_eps(z=z_center_slab, ax=ax2)
plt.show()

def post_process(sim_data):
    """Post-process simulation data: compute efficiency, RMSE, and plot diffraction results."""
    order = TARGET_ORDER
    number_of_orders = NUM_TARGET_ORDERS

    # Extract data
    plot_data = sim_data["diffractionmonitor"]
    intensity_measured = plot_data.power
    theta, phi = plot_data.angles
    theta = theta.isel(f=0)
    phi = phi.isel(f=0)
    power_values = plot_data.power.isel(f=0)
    total_power = 0

    # Calculate diffraction orders
    order1, power1, desiredPower1 = [], [], []

    for xorder in intensity_measured.orders_x:
        for yorder in intensity_measured.orders_y:
            val = (
                sim_data["diffractionmonitor"]
                .power.isel(f=0)
                .sel(orders_x=xorder, orders_y=yorder)
                .values
            )
            total_power += val
            if (abs(xorder) <= order) and (abs(yorder) <= order):
                order1.append((int(xorder), int(yorder)))
                power1.append(val)
                desiredPower1.append(1 / number_of_orders)

    rmse = np.sqrt(
        (1 / number_of_orders) * np.sum((np.array(power1) - np.sum(power1) / number_of_orders) ** 2)
    )
    rmse *= 100  # to get percentage

    labels1 = [f"({x},{y})" for x, y in order1]

    # Create a figure with two subplots side by side
    fig = plt.figure(figsize=(10, 4))
    ax_polar = fig.add_subplot(1, 2, 1, projection="polar")
    ax_bar = fig.add_subplot(1, 2, 2)

    # --- Polar plot ---
    sc = ax_polar.scatter(phi, theta, c=power_values, cmap="hot_r")
    fig.colorbar(sc, ax=ax_polar, orientation="vertical", pad=0.1)
    ax_polar.set_title("Far-Field Diffraction Pattern", va="bottom")

    # --- Bar plot ---
    total_power = float(np.sum(power_values))
    ax_bar.bar(
        range(len(power1)),
        100 * np.array(power1) / total_power,
        color="tab:blue",
        label="Measured (nonzero)",
    )
    ax_bar.bar(
        range(len(desiredPower1)),
        np.array(desiredPower1) * 100,
        color="cyan",
        alpha=0.3,
        label="Desired",
    )
    ax_bar.legend()
    ax_bar.set_xticks(range(len(power1)))
    ax_bar.set_xticklabels(labels1, rotation=90, fontsize=8)
    ax_bar.set_xlabel("Diffraction Order (x,y)")
    ax_bar.set_title("Diffraction Monitor Power by Order")

    plt.tight_layout()
    plt.show()

    efficiency = sum(power1) / total_power

    print(f"Efficiency: {efficiency:.2f}")
    print(f"RMSE: {rmse:.2f}")

    return efficiency, rmse

def loss_fn(params, beta, verbose=False):
    """Loss function for the design, the difference in intensity + the feature size penalty."""
    sim = make_sim(params, beta=beta)
    sim_data = web.run(sim, task_name="diffractive_beam_splitter", verbose=verbose)

    cost = intensity_diff_fn(sim_data)
    density = filter_project(params, beta)
    fab_penalty = erosion_dilation_penalty(density)

    res = cost**-1 - FAB_WEIGHT * fab_penalty

    # Extract per-order powers for monitoring (detached from autograd graph).
    power_da = sim_data["diffractionmonitor"].power.isel(f=0)
    pv = power_da.values
    orders_x = power_da.orders_x.values
    orders_y = power_da.orders_y.values

    order_powers = []
    order_labels = []
    target_sum = 0.0
    total_power = 0.0
    for i, ox in enumerate(orders_x):
        for j, oy in enumerate(orders_y):
            v = _to_float(pv[i, j])
            total_power += v
            if abs(ox) <= TARGET_ORDER and abs(oy) <= TARGET_ORDER:
                order_powers.append(v)
                order_labels.append((int(ox), int(oy)))
                target_sum += v

    aux_data = dict(
        objective=_to_float(cost),
        fab_penalty=_to_float(fab_penalty),
        efficiency=target_sum / total_power if total_power > 0 else 0.0,
        order_powers=np.array(order_powers),
        order_labels=order_labels,
    )

    return res, aux_data


loss_fn_val_grad = value_and_grad(loss_fn, has_aux=True)

(val, grad), aux_data = loss_fn_val_grad(params0, beta0, verbose=True)

print(f"Loss value (maximize): {val:.4f}")
print(f"  Cost (minimize): {aux_data['objective']:.4e}")
print(f"  Fab penalty: {aux_data['fab_penalty']:.4e}")
print(f"  Efficiency: {aux_data['efficiency']:.4f}")

print(f"Gradient norm: {np.linalg.norm(grad):.3e}")
print(f"Gradient range: [{np.min(grad):.3e}, {np.max(grad):.3e}]")

post_process(aux_data.get("sim_data", sim_data_norm))