"""Photonics2D problem.

This is essentially re-factored from the code at
https://nbviewer.org/github/fancompute/workshop-invdesign/blob/master/04_Invdes_wdm_scheduling.ipynb

Author: Mark Fuge @markfuge
"""

import dataclasses
from dataclasses import dataclass

# Need os import for makedirs for saving plots
import os
import pprint
from typing import Annotated, Any

# Importing autograd since the ceviche library uses it for automatic differentiation of the FDFD solver
import autograd.numpy as npa
from autograd.numpy.numpy_boxes import ArrayBox

# Import ArrayBox type for checking
import ceviche
from ceviche import fdfd_ez
from ceviche import jacobian
from ceviche.optimizers import adam_optimize
from gymnasium import spaces
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt

# --- EngiBench Core Imports ---
# Import necessary base classes and types directly from the library's core
from engibench.constraint import bounded
from engibench.constraint import IMPL
from engibench.constraint import THEORY
from engibench.core import ObjectiveDirection
from engibench.core import OptiStep
from engibench.core import Problem
from engibench.problems.photonics2d.backend import epsr_parameterization

# --- EngiBench Problem-Specific Backend ---
from engibench.problems.photonics2d.backend import init_domain
from engibench.problems.photonics2d.backend import insert_mode
from engibench.problems.photonics2d.backend import mode_overlap
from engibench.problems.photonics2d.backend import operator_blur
from engibench.problems.photonics2d.backend import operator_proj
from engibench.problems.photonics2d.backend import poly_ramp
from engibench.problems.photonics2d.backend import wavelength_to_frequency
from engibench.utils.upcast import upcast


class Photonics2D(Problem[npt.NDArray]):
    r"""Photonic Inverse Design 2D Problem (Wavelength Demultiplexer).

    ## Problem Description
    Optimize a 2D material distribution (`rho`) to function as a wavelength
    demultiplexer, routing wave with `lambda1` to output 1 and `lambda2` to output 2. The
    design variables represent material density that is converted to permittivity
    using filtering and projection.

    ### Motivation
    The optimization of photonic circuits in general, and multiplexers in particular, was
    one of the initial and most widely studied problems in the inverse design of
    electromagnetic/optical devices. In part, this is because multiplexer
    devices have several interesting properties that make them more difficult to create generative
    models of, compared to other problems in EngiBench. This includes the fact that, due to the wave
    properties of the physical phenomena there are usually multiple solutions with equivalent or
    similar performance, which results from shifting or inverting the phase profile of the
    electromagnetic wave. This adds complexity to the generative model in that the solution may not
    have a single unique global minimum. Another motivating factor for including this problem is the
    complexity of the structures/designs themselves: unlike in structural or thermal compliance
    problems, which lead to connected structures, the photonics solutions often involve several
    disconnected elements whose relative position and spacing is governed by the specific wavelengths
    it needs to demultiplex. This is a difficult prediction and generation task, compared to, e.g.,
    generating a connected beam structure. Thus, it acts as a good counterpoint to add to the library
    and provides a mechanism to benchmark generative algorithms that can perform well on both
    connected and disconnected design topologies.

    ## Design space
    This problem simulates a wavelength demultiplexer where the optimized device will direct an
    electromagnetic wave to one of two possible output ports depending on the wavelength/frequency of
    the incoming wave. Specifically, the demultiplexer targets two specific wavelengths (referred to
    $\lambda_1$ and $\lambda_2$ in the library), and the performance of the device is how well it can
    bend or direct the energy toward two specific locations in the device, as measured by how much of
    the electric field of each wavelength overlaps with the desired output port locations. The design
    space consists of a 2D array representing the presence of either a high or low permittivity,
    parameterized by `nelx` and `nely`, i.e., $\text{design_space} = [0,1]^{\text{nelx}\times \text{nely}}$.
    By default, the library uses a $120 \times 120$ space, however, this can be modified to non-square
    design spaces by the user. Specifically, we use a 2D tensor `rho` (num_elems_x, num_elems_y) with
    values in [0, 1], representing material density. This is stored as `design_space` (gymnasium.spaces.Box).

    ## Objectives
    The main objective is to maximize the overlap of the electric field of the simulated wavelength at
    the target output location, with an optional penalty for the amount of material used (this penalty
    weight is set to a small default value ($1e^{-2}$) for consistency, but can be altered for advanced
    usage):
    0. `total_overlap`: Objective to maximize, defined as
       `overlap1 * overlap2`. Higher is better. This corresponds to
       the overlap in the target electrical fields with the desired demultiplexing locations.
       Note that bot `simulate` and `optimize` subtract a small material penalty
       (`total_overlap - penalty`) to avoid multiple equivalent local optima, but this penalty
       is small relative to the overlap objective.

    ## Conditions
    These are designed as user-configurable parameters that alter the problem definition. The
    conditions consist of the two input wavelengths to be demultiplexed --
    $\lambda_1$ and $\lambda_2$, as well as a desired `blur_radius` ($r_{blur}$) parameter,
    which blurs (using a circular convolution) the pixelized design field for a chosen number of
    integer pixels\textemdash this blurring essentially creates a penalty on the minimum feature size
    of the design. The size of the device -- expressed as `nelx` and `nely` -- is also adjustable,
    and could be viewed as a possible condition for multi-resolution problems, but in practice, as with
    `Beams2D`, this is built into the problem definition since it produces a different dataset.

    Default problem parameters that can be overridden via the `config` dict:
    - `lambda1`: The first input wavelength in μm (default: 1.5 μm).
    - `lambda2`: The first input wavelength in μm (default: 1.3 μm).
    - `blur_radius`: Radius for the density blurring filter (default: 2).
                     Higher values correspond to larger elements, which could
                     possibly be more manufacturable.
    - `num_elems_x`: Number of grid cells in x (default: 120).
    - `num_elems_y`: Number of grid cells in y (default: 120).

    In practice, for the dataset loading, we will keep `num_elems_x` and `num_elems_y`to set
    values for each dataset, such that different resolutions correspond to different
    independent datasets.

    ## Optimization Parameters
    Note: These are advanced parameters that alter the optimization process --
    we do not recommend changing these if you are only using the library for benchmarking,
    as it could make results less reproducible across papers using this problem.)
    - `num_optimization_steps`: Total number of optimization steps (default: 300).
    - `step_size`: Adam optimizer step size (default: 1e-1).
    - `penalty_weight`: Weight for the L2 penalty term (default: 1e-2). Larger values reduce
                        unnecessary material, but may lead to worse performance if too large.
    - `eta`: Projection center parameter (default: 0.5). There is little reason to change this.
    - `N_proj`: Number of projection applications (default: 1). Increasing this can help make
                the design more binary.
    - `N_blur`: Number of blur applications (default: 1). Increasing this smooths the design more.
    - `initial_beta`: Initial beta for the optimization continuation scheme (default: 1.0).
    - `save_frame_interval`: Interval for saving intermediate design frames during optimization.
                             If > 0, saves a frame every `save_frame_interval` iterations
                             to the `opt_frames/` directory. Default is 0 (disabled).

    ## Internal Constants
    Note: These are not typically changed by users, but provided here for technical reference
    - `dl`: Spatial resolution (meters) (default: 40e-9).
    - `Npml`: Number of PML cells (default: 20).
    - `epsr_min`: Minimum relative permittivity (default: 1.0).
    - `epsr_max`: Maximum relative permittivity (default: 12.0).
    - `space_slice`: Extra space for source/probe slices (pixels) (default: 8).

    ## Simulator
    The simulation uses the `ceviche` library's Finite Difference Frequency Domain (FDFD)
    solver (`fdfd_ez`). Optimization uses `ceviche.optimizers.adam_optimize` with
    gradients computed via automatic differentiation (`autograd`).
    The simulation code uses the \texttt{ceviche} library and specifically,
    the wave demultiplexer demonstration case provided by the
    [library authors](https://github.com/fancompute/workshop-invdesign/blob/master/04_Invdes_wdm_scheduling.ipynb}
    based on their related publication, which uses a similar formalism to an earlier demultiplexer
    paper by Piggott (2015). The optimization method is first-order and uses the Adam optimizer. Beyond
    the baseline implementation already available via `ceviche`, we implemented a polynomial $\beta$ continuation
    scheme that performed more reliably than the step-wise continuation scheme used in the original implementation,
    and EngiBench also possesses the ability to change the starting and ending continuation values, for future
    research cases where one wishes to estimate or optimize the continuation schedule themselves. Other than these
    changes, the implementation of this problem is as consistent as possible with that of the original `ceviche` library.

    ## Dataset
    This problem offers a single datasets of `nelx`=120 and `nely`=120, although various sizes of `nelx` and `nely`
    could be generated from the library if desired. The dataset includes columns for the optimal design, all conditions
    listed above, and the corresponding objective value history as the optimizer optimized toward the optimal design
    provided in the dataset. The dataset was generated by sampling by sampling at random $\lambda_1$, $\lambda_2$ and
    $r_{blur}$ over the conditions mentioned above. The dataset is available on the
    [Hugging Face Datasets Hub](https://huggingface.co/datasets/IDEALLab/photonics_2d_120_120_v0).

    ### v0

    #### Fields
    Each dataset contains:
    - `lambda1`: The first input wavelength in μm.
    - `lambda2`: The second input wavelength in μm.
    - `blur_radius`: Radius for the density blurring filter (pixels).
    - `optimal_design`: The optimal design density array (shape num_elems_x, num_elems_y).
    - `optimization_history`: A list of objective values from the optimization process (field overlap  minus penalty, where higher is better) -- This is for advanced use.

    #### Creation Method
    To generate a dataset for training, we generate (randomly, uniformly) swept over the following parameters:
    - $\lambda_1 \in [0.5\mu m, 1.25\mu m]$ = `lambda1` = `rng.uniform(low=0.5, high=1.25, size=20)` -- This corresponds roughly to a portion of the visible spectrum up to near-infrared.
    - $\lambda_2 \in [0.75\mu m, 1.5\mu m]$ = `lambda2` = `rng.uniform(low=0.75, high=1.5, size=20)` -- This corresponds roughly to a portion of the visible spectrum up to near-infrared.
    - $r_{blur}$ = `blur_radius` = `range(0, 5)`

    ## Citation
    This problem is directly refactored from the Ceviche Library:
    https://github.com/fancompute/ceviche
    and if you use this problem your experiments, you can use the citation below
    provided by the original library authors:
    ```
    @article{hughes2019forward,
        title={Forward-Mode Differentiation of Maxwell's Equations},
        author={Hughes, Tyler W and Williamson, Ian AD and Minkov, Momchil and Fan, Shanhui},
        journal={ACS Photonics},
        volume={6},
        number={11},
        pages={3010--3016},
        year={2019},
        publisher={ACS Publications}
    }
    ```

    ## Lead
    Mark Fuge @markfuge
    """

    version = 0
    # --- Objective Definition ---
    objectives: tuple[tuple[str, ObjectiveDirection]] = (("total_overlap", ObjectiveDirection.MAXIMIZE),)
    # Note: there is also a small material penalty term added to the objective, but this is minor in comparison
    # We keep a single objective name for simplicity in the list.

    # Constants specific to problem design
    _pml_space = 10  # Space between PML and design region (pixels)
    _wg_width = 12  # Width of waveguides (pixels)
    _wg_shift = 9  # Lateral shift for output waveguides (pixels)
    _dl = 40e-9  # Spatial resolution (meters)
    _num_elems_pml = 20  # Number of PML cells (pixels)
    _epsr_min = 1.0  # Minimum relative permittivity (background)
    _epsr_max = 12.0  # Maximum relative permittivity (material)
    _space_slice = 8  # Extra space for source/probe slices (pixels)

    # Defaults for the optimization parameters
    _num_optimization_steps_default = 200  # Default number of optimization steps
    _step_size_default = 1e-1  # Default step size for Adam optimizer
    _eta_default = 0.5
    _num_projections_default = 1
    _penalty_weight_default = 1e-3  # Default weight for mass penalty term
    _num_blurs_default = 1
    _max_beta_default = 300  # Default maximum beta for scheduling
    _initial_beta_default = 1.0  # Default initial beta for scheduling

    @dataclass
    class Conditions:
        """Conditions."""

        lambda1: Annotated[
            float,
            bounded(lower=0.0).category(THEORY),
            bounded(lower=0.5).category(IMPL),
            bounded(lower=0.5, upper=1.5).warning().category(IMPL),
        ] = 1.5
        """First input wavelength in μm"""
        lambda2: Annotated[
            float,
            bounded(lower=0.0).category(THEORY),
            bounded(lower=0.5).category(IMPL),
            bounded(lower=0.5, upper=1.5).warning().category(IMPL),
        ] = 1.3
        """Second input wavelength in μm"""
        blur_radius: Annotated[
            int, bounded(lower=0).category(THEORY | IMPL), bounded(lower=0, upper=5).warning().category(IMPL)
        ] = 2
        """Radius for the density blurring filter (pixels)"""

    @dataclass
    class Config(Conditions):
        """Structured representation of configuration parameters for a numerical computation."""

        num_elems_x: Annotated[
            int,
            bounded(lower=0).category(THEORY),
            bounded(lower=60).category(IMPL),
            bounded(lower=90, upper=200).warning().category(IMPL),
        ] = 120
        """number of grid cells in x"""
        num_elems_y: Annotated[
            int,
            bounded(lower=0).category(THEORY),
            bounded(lower=105).category(IMPL),
            bounded(lower=110, upper=300).warning().category(IMPL),
        ] = 120
        """number of grid cells in y"""

    design_space = spaces.Box(low=0.0, high=1.0, shape=(Config.num_elems_x, Config.num_elems_y), dtype=np.float64)

    dataset_id = f"IDEALLab/photonics_2d_{Config.num_elems_x}_{Config.num_elems_y}_v0"
    container_id = None

    def __init__(
        self,
        seed: int = 0,
        config: dict[str, Any] | None = None,
        num_elems_x: int = Config.num_elems_x,
        num_elems_y: int = Config.num_elems_y,
        **kwargs,
    ) -> None:
        """Initializes the Photonics2D problem.

        Args:
            seed (int): The random seed for the problem.
            config (dict): A dictionary with configuration (e.g., boundary conditions) for the simulation.
            num_elems_x (int): Number of grid cells in x (default: 120).
            num_elems_y (int): Number of grid cells in y (default: 120).
            **kwargs: Additional keyword arguments.
        """
        super().__init__(seed=seed, **kwargs)

        # Replace the conditions with any new configs passed in
        self.config = self.Config(num_elems_x=num_elems_x, num_elems_y=num_elems_y, **(config or {}))
        # Restrict config to conditions:
        self.conditions = upcast(self.config)

        print("Initializing Photonics Problem with configuration:")
        pprint.pp(self.conditions)
        self.num_elems_x = num_elems_x
        self.num_elems_y = num_elems_y
        self.design_space = spaces.Box(low=0.0, high=1.0, shape=(num_elems_x, num_elems_y), dtype=np.float32)
        self.dataset_id = f"IDEALLab/photonics_2d_{num_elems_x}_{num_elems_y}_v{self.version}"

        # Setup basic simulation parameters
        self.omega1 = wavelength_to_frequency(self.conditions.lambda1)
        self.omega2 = wavelength_to_frequency(self.conditions.lambda2)
        self._current_beta: float = self._max_beta_default

    def _setup_simulation(self, config: dict[str, Any] | None = None) -> dict[str, Any]:
        """Helper function to setup simulation parameters and domain."""
        # Merge config with default conditions
        current_conditions = {**dataclasses.asdict(self.conditions), **(config or {})}

        # Initialize domain geometry
        self._bg_rho, self._design_region, self._input_slice, self._output_slice1, self._output_slice2 = init_domain(
            num_elems_x=self.num_elems_x,
            num_elems_y=self.num_elems_y,
            num_elems_pml=self._num_elems_pml,
            space=self._pml_space,
            wg_width=self._wg_width,
            wg_shift=self._wg_shift,
            space_slice=self._space_slice,
        )

        return current_conditions

    def _run_fdfd(
        self, design: npt.NDArray, conditions: dict[str, Any]
    ) -> tuple[npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray]:
        """Helper to run FDFD and return key components (epsr, fields, sources, probes)."""
        omega1 = self.omega1
        omega2 = self.omega2
        beta = self._current_beta
        num_blurs = conditions.get("num_blurs", self._num_blurs_default)
        num_projections = conditions.get("num_projections", self._num_projections_default)
        eta = conditions.get("eta", self._eta_default)

        # 1. Parameterize
        epsr = epsr_parameterization(
            rho=design,
            bg_rho=self._bg_rho,
            design_region=self._design_region,
            radius=conditions["blur_radius"],
            num_blurs=num_blurs,
            beta=beta,
            eta=eta,
            num_projections=num_projections,
            epsr_min=self._epsr_min,
            epsr_max=self._epsr_max,
        )

        # 2. Setup Sources and Probes (depend on epsr)
        source1 = insert_mode(omega1, self._dl, self._input_slice.x, self._input_slice.y, epsr, m=1)
        source2 = insert_mode(omega2, self._dl, self._input_slice.x, self._input_slice.y, epsr, m=1)
        probe1 = insert_mode(omega1, self._dl, self._output_slice1.x, self._output_slice1.y, epsr, m=1)
        probe2 = insert_mode(omega2, self._dl, self._output_slice2.x, self._output_slice2.y, epsr, m=1)

        # 3. Setup FDFD Simulations
        # We need to run two simulations, one for each wavelength, to see which paths the light takes in both
        self._simulation1 = fdfd_ez(omega1, self._dl, epsr, [self._num_elems_pml, self._num_elems_pml])
        self._simulation2 = fdfd_ez(omega2, self._dl, epsr, [self._num_elems_pml, self._num_elems_pml])

        # 4. Solve FDFD
        # Always solve as Ez fields are needed for return value or objective calc
        _, _, ez1 = self._simulation1.solve(source1)
        _, _, ez2 = self._simulation2.solve(source2)

        return epsr, ez1, ez2, source1, source2, probe1, probe2

    def simulate(self, design: npt.NDArray, config: dict[str, Any] | None = None, **kwargs) -> npt.NDArray:  # noqa: ARG002
        """Simulates the performance of a design, returning the raw objective value.

           Stores simulation fields (`Ez1`, `Ez2`, `epsr`) internally in `_last_Ez1`,
           `_last_Ez2`, `_last_epsr` for later access (e.g., by render).

        Args:
            design (npt.NDArray): The design density array `rho` (shape num_elems_x, num_elems_y).
            config (dict): Dictionary to override default conditions.
            **kwargs: Additional keyword arguments (ignored).

        Returns:
            npt.NDArray: 1-element array: [total_overlap - penalty], where higher is better.
        """
        conditions = self._setup_simulation(config)

        # --- Run Simulation ---
        # We don't need source returns here
        print("Simulating design under the following conditions:")
        pprint.pp(conditions)
        epsr, ez1, ez2, _, _, probe1, probe2 = self._run_fdfd(design, conditions)

        # --- Store Results Internally ---
        self._last_epsr = epsr.copy()
        self._last_Ez1 = ez1.copy()
        self._last_Ez2 = ez2.copy()

        # --- Calculate Objective ---
        # Use standard numpy here, no gradients needed
        overlap1 = np.abs(np.sum(np.conj(ez1) * probe1)) * 1e6
        overlap2 = np.abs(np.sum(np.conj(ez2) * probe2)) * 1e6
        total_overlap = overlap1 * overlap2  # Maximize this
        penalty_weight = conditions.get("penalty_weight", self._penalty_weight_default)
        penalty = penalty_weight * np.linalg.norm(design)

        return np.array([total_overlap - penalty], dtype=np.float64)

    def optimize(  # noqa: PLR0915
        self,
        starting_point: npt.NDArray,
        config: dict[str, Any] | None = None,
        **kwargs,  # noqa: ARG002
    ) -> tuple[npt.NDArray, list[OptiStep]]:
        """Optimizes a topology (rho) starting from `starting_point` using Adam.

           Maximizes `total_overlap - penalty` using gradients from autograd via Ceviche.
           Optionally saves intermediate design frames based on `save_frame_interval`.

        Args:
            starting_point (npt.NDArray): The starting design `rho` (shape num_elems_x, num_elems_y).
            config (dict): Dictionary to override default conditions (e.g., num_optimization_steps,
                           step_size, penalty_weight, save_frame_interval).
            **kwargs: Additional keyword arguments (ignored).

        Returns:
            tuple[npt.NDArray, list[OptiStep]]:
                - The optimized design `rho` (float32, shape num_elems_x, num_elems_y).
                - A list of OptiStep history. `OptiStep.obj_values` contains the
                  value of the internally optimized objective (i.e., `total_overlap - penalty`).
                  The `step` attribute corresponds to the optimizer iteration.
        """
        conditions = self._setup_simulation(config)

        # Reset the current beta to one for the optimization
        initial_beta = conditions.get("initial_beta", self._initial_beta_default)
        # Set current beta to the max so that first opt history is under simulate conditions
        # Later in optimize we will use initial_beta to perform continuation
        self._current_beta = self._max_beta_default

        print("Attempting to run Optimization for Photonics2D under the following conditions:")
        pprint.pp(conditions)

        # Pull out problem-specific parameters from conditions
        num_elems_x = self.num_elems_x
        num_elems_y = self.num_elems_y
        # Pull out optimization parameters from conditions
        # Parameters specific to optimization
        num_optimization_steps = conditions.get("num_optimization_steps", self._num_optimization_steps_default)
        step_size = conditions.get("step_size", self._step_size_default)
        penalty_weight = conditions.get("penalty_weight", self._penalty_weight_default)
        self._eta = conditions.get("eta", self._eta_default)
        self._num_projections = conditions.get("num_projections", self._num_projections_default)
        self._num_blurs = conditions.get("num_blurs", self._num_blurs_default)
        # --- Get the frame saving interval from conditions for plotting ---
        save_frame_interval = conditions.get("save_frame_interval", 0)

        # --- Initial Simulation for first OptiStep history (IOG calculation) ---
        print("Optimize: Calculating initial design...")  # Keep this info message
        _epsr_init, ez1_init, ez2_init, source1_init, source2_init, probe1_init, probe2_init = self._run_fdfd(
            starting_point, conditions
        )

        self._source1 = source1_init
        self._source2 = source2_init
        self._probe1 = probe1_init
        self._probe2 = probe2_init
        # Calculate overlaps
        initial_overlap1 = mode_overlap(ez1_init, self._probe1)
        initial_overlap2 = mode_overlap(ez2_init, self._probe2)
        total_overlap_initial = initial_overlap1 * initial_overlap2
        # Calculate initial material penalty
        initial_penalty = penalty_weight * np.linalg.norm(starting_point)

        # Add initial design performance into opti_steps_history
        opti_steps_history: list[OptiStep] = []
        initial_design_optistep = OptiStep(
            obj_values=np.array([total_overlap_initial - initial_penalty], dtype=np.float64), step=0
        )
        opti_steps_history.append(initial_design_optistep)

        # Ensure directory exists for saving frames
        frame_dir = "opt_frames"
        os.makedirs(frame_dir, exist_ok=True)

        # --- Define Objective Function for Ceviche Optimizer ---
        def objective_for_optimizer(rho_flat: npt.NDArray | ArrayBox) -> float | ArrayBox:
            """Calculates (overlap - penalty) for maximization.

            Note: All functions or inputs here should be compatible with autograd (npa).
            """
            rho = rho_flat.reshape((num_elems_x, num_elems_y))

            # --- Parameterization and Simulation ---
            epsr = epsr_parameterization(
                rho=rho,
                bg_rho=self._bg_rho,
                design_region=self._design_region,
                radius=conditions["blur_radius"],
                num_blurs=self._num_blurs,
                beta=self._current_beta,
                eta=self._eta,
                num_projections=self._num_projections,
                epsr_min=self._epsr_min,
                epsr_max=self._epsr_max,
            )
            self._simulation1.eps_r = epsr
            self._simulation2.eps_r = epsr
            _, _, ez1 = self._simulation1.solve(self._source1)
            _, _, ez2 = self._simulation2.solve(self._source2)

            # Calculate overlaps
            overlap1 = mode_overlap(ez1, self._probe1)
            overlap2 = mode_overlap(ez2, self._probe2)
            total_overlap = overlap1 * overlap2
            penalty = penalty_weight * npa.linalg.norm(rho)
            return total_overlap - penalty  # Value to MAXIMIZE

        # --- Define Gradient ---
        objective_jac = jacobian(objective_for_optimizer, mode="reverse")

        # --- Define Callback ---
        def callback(iteration: int, objective_history_list: list, rho_flat: npt.NDArray | ArrayBox) -> None:
            """Callback for adam_optimize. Receives the history of objective values."""
            # Handle Empty History
            if not objective_history_list:
                return

            # Get the latest objective value
            last_scalar_obj_value = objective_history_list[-1]

            # Beta Scheduling Logic -- Quadratic ramp from 0 to max_beta
            iteration = len(objective_history_list)
            self._current_beta = poly_ramp(
                iteration, max_iter=num_optimization_steps, b0=initial_beta, bmax=self._max_beta_default, degree=2
            )

            # Store OptiStep info
            step_info = OptiStep(obj_values=np.array([last_scalar_obj_value], dtype=np.float64), step=iteration)
            opti_steps_history.append(step_info)

            # --- Configurable Intermediate Frame Saving ---
            # Check if saving is enabled and if current iteration is a multiple of the interval
            # Also check iteration > 0 to avoid saving the initial state redundantly
            if (
                save_frame_interval is not None
                and save_frame_interval > 0
                and iteration > 0
                and iteration % save_frame_interval == 0
            ):
                # Reshape the current design parameters
                current_rho = rho_flat.reshape((num_elems_x, num_elems_y))
                current_rho = operator_proj(current_rho, self._eta, beta=self._current_beta, num_projections=1)

                # --- Call self.render to generate the plot ---
                # Pass the current conditions dictionary in case render needs it
                # Note: This will re-run the simulation for the current_rho
                fig = self.render(current_rho, open_window=False, config=conditions)
                # ---------------------------------------------

                save_path = os.path.join(frame_dir, f"frame_iter_{iteration:04d}.png")

                # Save the figure returned by render
                fig.savefig(save_path, dpi=200)
                plt.close(fig)  # Close the figure to free memory
                print(f"Callback Iter {iteration}: Saved frame to {save_path}")
            # --- End Frame Saving ---
            if iteration == num_optimization_steps - 1:
                print(f"Final Iteration {iteration}: Objective Value: {last_scalar_obj_value:.3e}")
                print("Saving render of final design...")
                current_rho = rho_flat.reshape((num_elems_x, num_elems_y))
                current_rho = operator_proj(current_rho, self._eta, beta=self._current_beta, num_projections=1)
                current_rho = np.rint(current_rho).astype(np.float32)  # Convert to binary for final save
                fig = self.render(current_rho, open_window=False, config=conditions)
                save_path = os.path.join(frame_dir, "frame_final.png")
                fig.savefig(save_path, dpi=200)
                plt.close(fig)

        # --- Run Optimization ---
        print(
            f"\nStarting optimization with num_optimization_steps={num_optimization_steps}, step_size={step_size}"
        )  # Keep start message
        (rho_optimum_flat, _) = adam_optimize(
            objective_for_optimizer,
            starting_point.flatten(),
            objective_jac,
            Nsteps=num_optimization_steps,
            direction="max",
            step_size=step_size,
            callback=callback,
        )

        # --- Final Result ---
        rho_optimum = rho_optimum_flat.reshape((num_elems_x, num_elems_y))
        # Project the optimized design to the valid range [0, 1]
        rho_optimum = operator_proj(rho_optimum, self._eta, beta=self._current_beta, num_projections=1)
        rho_optimum = np.rint(rho_optimum)  # Convert to binary for final
        return rho_optimum.astype(np.float32), opti_steps_history

    # --- render method remains the same as previous version ---
    def render(self, design: npt.NDArray, config: dict[str, Any] | None = None, *, open_window: bool = False) -> Figure:
        """Renders the design (rho) and the resulting E-field magnitudes.

           Runs a simulation for the provided design to get the fields for plotting.
           Uses the internally stored fields if called immediately after `simulate`.

        Args:
            design (npt.NDArray): The design `rho` to render.
            open_window (bool): If True, opens a window with the rendered plot.
            config (dict): Config overrides for simulation parameters if needed for rendering.
            **kwargs: Additional keyword arguments (ignored).


        Returns:
            Figure: The matplotlib Figure object containing three plots:
                        |Ez| at omega1, |Ez| at omega2, and Permittivity (eps_r).
        """
        conditions = self._setup_simulation(config)

        print("Rendering design under the following conditions:")
        pprint.pp(conditions)
        # Use run_fdfd but ignore most outputs, just need epsr, ez1, ez2
        epsr, ez1, ez2, _source1, _source2, probe1, probe2 = self._run_fdfd(design, conditions)

        overlap1 = mode_overlap(ez1, probe1)
        overlap2 = mode_overlap(ez2, probe2)
        total_overlap = overlap1 * overlap2

        # Store these fields as the "last" simulated ones as well
        self._last_epsr = epsr.copy()
        self._last_Ez1 = ez1.copy()
        self._last_Ez2 = ez2.copy()

        # --- Plotting ---
        fig, ax = plt.subplots(1, 3, constrained_layout=True, figsize=(9, 3))
        ceviche.viz.abs(
            ez1,
            outline=epsr,
            ax=ax[0],
            cbar=False,
            outline_alpha=0.25,
        )
        ceviche.viz.abs(
            ez2,
            outline=epsr,
            ax=ax[1],
            cbar=False,
            outline_alpha=0.25,
        )
        ceviche.viz.real(epsr, ax=ax[2], cmap="Greys")
        slices_to_plot = [self._input_slice, self._output_slice1, self._output_slice2]
        for sl in slices_to_plot:
            if sl:
                for axis in ax[:2]:  # Plot on field plots
                    axis.plot(sl.x * np.ones(len(sl.y)), sl.y, "w-", alpha=0.5, linewidth=1)
        lambda1_um = conditions["lambda1"]
        lambda2_um = conditions["lambda2"]
        blur_radius = conditions["blur_radius"]
        fig.suptitle(
            f"Total Overlap: {total_overlap:.2f} (λ1={lambda1_um:.2f} μm, λ2={lambda2_um:.2f} μm, blur={blur_radius}, $\\beta$ = {self._current_beta:.2f})",
        )
        ax[0].set_title(f"|Ez| at $\\lambda_1$ = {lambda1_um:.2f} $\\mu$m")
        ax[1].set_title(f"|Ez| at $\\lambda_2$ = {lambda2_um:.2f} $\\mu$m")
        ax[2].set_title(r"Permittivity $\epsilon_r$")
        for axis in ax:
            axis.set_xlabel("")
            axis.set_ylabel("")
            axis.set_xticks([])
            axis.set_yticks([])

        plt.tight_layout()
        if open_window:
            plt.show(block=False)
        return fig

    def _randomized_noise_field_design(self, noise: float = 0.001, blur: int = 0) -> npt.NDArray:
        """Generates a starting design with small random variations.

           Creates a design that is 0.5 within the design region, plus small
           normal random noise (0.001 * randn). Returns 0 as the index placeholder.

        Args:
            noise (float): The amount of noise to add to the uniform field.
            blur (float): The amount of blurring to apply to random field. if >0 can produce
                          different local optima for Adam, and thus can be useful for
                          exploring multiple local optima in the problem. Disabled by default.

        Returns:
            tuple[npt.NDArray, int]: The starting design array (rho) and an integer (0).
        """
        space = self._pml_space
        num_elems_pml = self._num_elems_pml

        design_region = np.zeros((self.num_elems_x, self.num_elems_y))
        design_region[
            num_elems_pml + space : self.num_elems_x - num_elems_pml - space,
            num_elems_pml + space : self.num_elems_y - num_elems_pml - space,
        ] = 1

        # Ensure np_random is initialized
        if self.np_random is None:
            self.reset()
        # Generate random numbers using the problem's RNG
        # Use randomized initialization -- for now keep
        random_noise = noise * self.np_random.standard_normal((self.num_elems_x, self.num_elems_y))
        rho_start = design_region * (0.5 + random_noise)
        if blur > 0.0:
            rho_start = operator_blur(rho_start, blur)

        return rho_start.astype(np.float32)

    def random_design(
        self, noise: float | None = None, blur: int = 0, dataset_split: str = "train", design_key: str = "optimal_design"
    ) -> tuple[npt.NDArray, int]:
        """Generates a random initial design.

        Can return a design with small random variations or a uniform design, or can pull
        from the datasets (when available).

        Args:
            noise (float|None): If None, pull from dataset. If float, use that as the noise level.
            blur (int): The amount of pixel blurring to apply to random field. Only active if noise is used.
            dataset_split (str): The key for the dataset to sample from.
            design_key (str): The key for the design to sample from.

        Returns:
            tuple[npt.NDArray, int]: The starting design array (rho) and an integer (0).
        """
        # Ensure np_random is initialized
        if self.np_random is None:
            self.reset()

        if noise is not None:
            rho_start = self._randomized_noise_field_design(noise=noise, blur=blur)
            return rho_start, 0
        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split]), dtype=int)
        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd

    def reset(self, seed: int | None = None, **kwargs) -> None:
        """Resets the problem, which in this case, is just the random seed."""
        return super().reset(seed, **kwargs)


# --- Example Usage (main block) ---
if __name__ == "__main__":
    # Problem Configuration Example
    problem_config = {
        "lambda1": 1.07,
        "lambda2": 0.84,
        "blur_radius": 1,
    }
    problem = Photonics2D(config=problem_config, num_elems_x=120, num_elems_y=120, seed=42)

    start_design, _ = problem.random_design(noise=0.1, blur=1)  # Randomized design with noise
    fig_start = problem.render(start_design)

    # Simulation Example
    print("Simulating starting design...")
    # Simulate returns the raw objective = penalty - overlap1*overlap2
    obj_start_raw = problem.simulate(start_design)
    print(f"Starting Raw Objective ({problem.objectives_keys[0]}): {obj_start_raw[0]:.4f}")

    # Optimization Example
    # Advanced Usage: Modifying optimization parameters
    opt_config = {"num_optimization_steps": 200, "save_frame_interval": 2, "initial_beta": 1.0}
    problem.reset(seed=42)
    print(f"Optimizing design with ({opt_config})...")
    # Optimize maximizes (overlap - penalty)
    optimized_design, opti_history = problem.optimize(start_design, config=opt_config)
    print(f"Optimization finished. History length: {len(opti_history)}")
    if opti_history:
        print(f"First step objective: {opti_history[0].obj_values[0]:.4f}")
        print(f"Last step objective: {opti_history[-1].obj_values[0]:.4f}")

    print("Rendering optimized design...")
    fig_opt = problem.render(optimized_design, open_window=True)
    frame_dir = "opt_frames"
    fig_opt.savefig(frame_dir + "/optimized_design.png", dpi=200)

    print("Simulating the final optimized design...")
    # Simulate returns the raw objective = penalty - overlap1*overlap2
    problem.reset(seed=43)
    obj_opt_raw = problem.simulate(optimized_design)
    print(f"Optimized Raw Objective ({problem.objectives_keys[0]}): {obj_opt_raw[0]:.4f}")

    if plt.get_fignums():
        print("Close plot window(s) to exit.")
        plt.show()
