from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import gym
import numpy as np
import plotly.graph_objects as go
import torch
from plotly.basedatatypes import BaseTraceType
from skfem import Mesh
from torch_geometric.data import Data

from .fem_problem import FEMProblemCircularQueue, FEMProblemWrapper
from .fem_util import (
    construct_sizing_field_1d,
    get_aggregation_per_element,
    get_triangle_areas_from_indices,
    sample_in_range,
)
from .utils import save_concatenate
from .visualization import get_plotly_mesh_traces_and_layout

class MeshRefinement(gym.Env):
    """Graph-based 2D mesh refinement RL environment using scikit-FEM backend."""

    def __init__(
        self, environment_config: Dict[Union[str, int], Any], seed: Optional[int] = None
    ):
        """
        Args:
            environment_config: Config for the environment.
                Details can be found in the configs/references/mesh_refinement_reference.yaml example file
            seed: Optional seed for the random number generator.
        """
        self._environment_config = environment_config
        self._random_state: np.random.RandomState = np.random.RandomState(seed=seed)
        self._num_node_features: int = 0
        self._num_edge_features: int = 0

        self.fem_problem_queue = FEMProblemCircularQueue(
            fem_config=environment_config.get("fem"),
            random_state=np.random.RandomState(seed=seed),
        )
        self.fem_problem: Optional[FEMProblemWrapper] = None

        ################################################
        #        general environment parameters        #
        ################################################
        self._refinement_strategy: str = environment_config.get("refinement_strategy")
        self._max_timesteps = environment_config.get("num_timesteps")
        self._element_limit_penalty = environment_config.get("element_limit_penalty")
        self._maximum_elements = environment_config.get("maximum_elements")
        self._element_penalty_config = self._environment_config.get("element_penalty")
        self._sample_penalty = self._element_penalty_config.get("sample_penalty")

        ################################################
        # graph connectivity, feature and action space #
        ################################################
        self._reward_type = environment_config.get("reward_type")
        self._global_reward_alpha = float(environment_config.get("global_reward_alpha", 0.2))
        _rho_w = environment_config.get("rho_weights", {})
        self._w_rho_int = float(_rho_w.get("w_int", 1.0))
        self._w_rho_jump = float(_rho_w.get("w_jump", 1.0))
        self._w_rho_sbc = float(_rho_w.get("w_sbc", 1.0))
        self._include_vertices = environment_config.get("include_vertices")
        self._set_graph_sizes()

        ################################################
        #          internal state and cache            #
        ################################################
        self._timestep: int = 0
        self._element_penalty_lambda = None  # 0  # set default value
        self._initial_approximation_errors: Optional[Dict[str, float]] = None
        self._reward = None
        self._cumulative_return: np.array = 0  # return of the environment
        # dictionary containing the error estimation for the current solution for different error evaluation metrics
        self._error_estimation_dict: Optional[Dict[str, np.array]] = None
        self._initial_error_norm = None
        self.current_error = None  # scalar total error, updated each step
        self.initial_error = (
            None  # scalar initial total error, fixed as normalization baseline
        )

        # last-step history for delta-based rewards and plotting
        self._previous_error_per_element: Optional[np.array] = None
        self._previous_num_elements: Optional[int] = None
        self._previous_agent_mapping = None
        self._previous_element_volumes = None
        self._previous_std_per_element = None
        self._previous_eta_components: Optional[Dict[str, np.ndarray]] = None
        self._previous_rho_components: Optional[Dict[str, np.ndarray]] = None

        # fields/internal variables for spatial mesh refinement, especially a spatial reward
        self._agent_mapping = None  # mapping List[old_element_indices] of size new_element_indices that maps
        self._reward_per_agent: Optional[np.array] = (
            0  # cumulative return of the environment per agent
        )
        self._cumulative_reward_per_agent: Optional[np.array] = (
            0  # cumulative reward of the environment per agent
        )

        # additional policy information that is not passed through the graph
        self._include_additional_policy_information = environment_config.get(
            "include_additional_policy_information"
        )

        self._manual_normalization = environment_config.get(
            "manual_normalization", None
        )  # manually normalize the error
        ################################################
        #            recording and plotting            #
        ################################################
        self._initial_num_elements = None

    def _set_graph_sizes(self):
        """
        Internally sets the
        * action dimension
        * number of node types and node features for each type
        * number of edge types and edge features for each type
        depending on the configuration. Uses the same edge features for all edge types.
        Returns:

        """
        edge_feature_config = self._environment_config.get("edge_features")
        self._edge_features = [
            feature_name
            for feature_name, include_feature in edge_feature_config.items()
            if include_feature
        ]
        # set number of edge features
        num_edge_features = 0
        if "euclidean_distance" in self._edge_features:
            num_edge_features += 1

        self._element_feature_functions = self._register_element_features()

        self._num_node_features = len(self._element_feature_functions)
        self._num_node_features += self.fem_problem_queue.num_pde_element_features
        self._num_edge_features = num_edge_features

    def _register_element_features(self) -> Dict[str, Callable[[], np.array]]:
        cfg = self._environment_config.get("element_features")
        names = [n for n, inc in cfg.items() if inc]
        feats = {}

        if "x_position" in names:
            feats["x_position"] = lambda: self._element_midpoints[:, 0]
        if "y_position" in names:
            feats["y_position"] = lambda: self._element_midpoints[:, 1]
        if "volume" in names:
            feats["volume"] = lambda: self._volume_normalized
        if "solution_std" in names:
            feats["internal_residual"] = lambda: self._residual_components["internal_residual"]
            feats["gradient_jump"] = lambda: self._residual_components["gradient_jump"]
            feats["sbc_residual"] = lambda: self._residual_components["sbc_residual"]
        if "element_penalty" in names:
            feats["element_penalty"] = lambda: np.repeat(self._element_penalty_lambda, self._num_elements)
        if "timestep" in names:
            feats["timestep"] = lambda: np.repeat(self._timestep, self._num_elements)
        if "wave_number" in names:
            feats["wave_number"] = lambda: np.repeat(self._wave_number, self._num_elements)
        if "k_local_sqrt_vol" in names:
            feats["k_local_sqrt_vol"] = lambda: self._k_local_sqrt_vol
        if "is_sbc_boundary" in names:
            feats["is_sbc_boundary"] = lambda: self._residual_components["is_sbc_boundary"]
        if "dist_to_interface" in names:
            feats["dist_to_interface"] = lambda: self._dist_to_interface

        # Complex field decomposition (always present for Helmholtz)
        feats["epsilon_r"] = lambda: self._epsilon_r_elements
        feats["total_solution_magnitude"] = lambda: np.abs(self._complex_solution_mean)
        return feats

    def reset(self) -> Data:
        """
        Resets the environment and returns an (initial) observation of the next rollout
        according to the reset environment state

        Returns:
            The observation of the initial state.
        """
        # get the next fem problem. This samples a new domain and new load function, resets the mesh and the solution.
        self.fem_problem = self.fem_problem_queue.next()

        # calculate the solution of the finite element problem for the initial mesh and retrieve an error per element
        self._error_estimation_dict = (
            self.fem_problem.calculate_solution_and_get_error()
        )

        # reset the internal state of the environment. This includes the current timestep, the current element penalty
        # and some values for calculating the reward and plotting the env
        self._reset_internal_state()

        observation = self.last_observation
        return observation

    def _reset_internal_state(self):
        """
        Resets the internal state of the environment
        Returns:

        """
        self._agent_mapping = np.arange(self._num_elements).astype(
            np.int64
        )  # map to identity at first step
        self._previous_agent_mapping = np.arange(self._num_elements).astype(
            np.int64
        )  # map to identity at first step
        self._previous_element_volumes = self.element_volumes
        self._previous_eta_indicator = self._eta_indicator
        self._previous_eta_components = self._eta_components_raw
        self._previous_rho_components = self._rho_components
        self._previous_solution_l2_norm = self._compute_solution_l2_norm()
        self._reward_per_agent = np.zeros(self.num_agents)
        self._cumulative_reward_per_agent = np.zeros(self._num_elements)

        # reset timestep and rewards
        self._timestep = 0
        self._reward = 0
        self._cumulative_return = 0
        self._diag_selected_count = -1   # 防止跨 episode 残留触发 is_terminal

        # reset internal state that tracks statistics over the episode
        self._previous_error_per_element = self.error_per_element

        # collect a dictionary of initial errors to normalize them when calculating metrics during evaluation
        self._initial_approximation_errors = (
            self._calculate_initial_approximation_errors()
        )

        self._previous_num_elements = self._num_elements
        self._initial_num_elements = self._num_elements

        self._initial_median_area = float(np.median(self.element_volumes))

        k = self._wave_number
        eps_r_elem = self._epsilon_r_elements
        lambda_local = 2.0 * np.pi / (k * np.sqrt(np.maximum(eps_r_elem, 1.0)))
        A_budget = 0.5 * (lambda_local / 6.0) ** 2
        self._element_budget_area = A_budget
        N_phys = int(np.ceil(np.sum(self.element_volumes / A_budget)))
        rho_min = 5.0
        self._n_budget = max(N_phys, int(np.ceil(rho_min * self._num_elements)))

        if self.error_per_element is not None:
            self._initial_error_norm = np.linalg.norm(self.error_per_element, axis=0)
            # Record initial total error as normalization baseline for reward calculation
            self.current_error = self._compute_total_error()
            self.initial_error = self.current_error
            # Protection against near-zero initial error (prevents division by zero)
            if self.initial_error < 1e-8:
                self.initial_error = 1.0

        # reset the element penalty, necessary if it is sampled
        if self._sample_penalty:
            sampling_type = self._element_penalty_config.get(
                "sampling_type", "loguniform"
            )
            min_value = self._element_penalty_config.get("min")
            max_value = self._element_penalty_config.get("max")
            element_penalty_lambda = sample_in_range(
                max_value, min_value, sampling_type
            )
            self._element_penalty_lambda = element_penalty_lambda
        else:  # element penalty is a scalar value
            self._element_penalty_lambda = self._element_penalty_config.get("value")

    def _calculate_initial_approximation_errors(self):
        if self._manual_normalization:
            return {
                error_name: self._manual_normalization
                for error_name in self.error_estimation_dict
            }
        else:
            result = {}
            for error_name, errors in self.error_estimation_dict.items():
                errors = np.atleast_1d(np.asarray(errors, dtype=np.float64))
                val = np.sqrt(np.sum(errors ** 2))
                result[error_name] = float(val) + 1e-12
            return result

    def step(self, action: np.ndarray) -> Tuple[Data, np.array, bool, Dict[str, Any]]:
        """
        Performs a step of the Mesh Refinement task.

        Wrapped in try-except to prevent program crashes from ill-conditioned
        FEM solves caused by degenerate meshes (especially in early training
        when the continuous sizing field produces extreme element shapes).

        On FEM failure: returns done=True with an extreme penalty reward (-10000)
        to implicitly teach the agent to avoid generating invalid meshes.

        Args:
            action: the action the agents will take in this step. Has shape (num_agents, action_dimension)
            Given as an array of shape (num_agents, action_dimension)

        Returns: A 4-tuple (observations, reward, done, info), where
            * observations is a graph of the agents and their positions, in this case of the refined mesh
            * reward is a single scalar shared between all agents, i.e., per **graph**
            * done is a boolean flag that says whether the current rollout is finished or not
            * info is a dictionary containing additional information
        """
        assert not self.is_terminal, (
            f"Tried to perform a step on a terminated environment. Currently on step "
            f"{self._timestep:} of {self._max_timesteps:} "
            f"with {self._num_elements}/{self._maximum_elements} elements."
        )

        self._timestep += 1

        # ================================================================
        #  核心逻辑: try-except 物理防崩盘机制
        #  捕获 FEM 求解器因畸形网格抛出的任何异常
        # ================================================================
        try:
            self._set_previous_step()

            # refine mesh and store which element has become which set of new elements
            self._agent_mapping = self._refine_mesh(action=action)

            # solve equation and calculate error per element/element
            self._previous_error_per_element = self.error_per_element

            self._error_estimation_dict = (
                self.fem_problem.calculate_solution_and_get_error()
            )

            # query returns
            observation = self.last_observation

            reward_dict = self._get_reward_dict()
            metric_dict = self._get_metric_dict()
            action_dict = self._get_action_dict(action=action)

            # done after a given number of steps or if the mesh becomes too large
            done = self.is_terminal
            info = (
                reward_dict
                | metric_dict
                | action_dict
                | {
                    "is_truncated": self.is_truncated,
                    "return": self._cumulative_return,
                    "neg_action_ratio": getattr(self, "_diag_neg_ratio", 0.0),
                    "eligible_ratio": getattr(self, "_diag_eligible_ratio", 0.0),
                    "masked_ratio": getattr(self, "_diag_masked_ratio", 0.0),
                    "selected_count": getattr(self, "_diag_selected_count", 0),
                    "n_budget": self._n_budget,
                }
            )
            return observation, self._reward, done, info

        # except (np.linalg.LinAlgError, ValueError, RuntimeError, Exception) as e:
        except (np.linalg.LinAlgError, ValueError, RuntimeError) as e:
            # ============================================================
            #  FEM 物理崩溃捕获
            #  可能原因:
            #    1. 畸形网格导致刚度矩阵奇异 (LinAlgError)
            #    2. 连续动作产生了退化元素 (ValueError)
            #    3. scikit-fem 内部网格操作异常 (RuntimeError)
            #
            #  策略: 立即终止本回合，给予极端惩罚，迫使智能体学习
            #       避免产生会导致 FEM 崩溃的网格。
            # ============================================================
            import sys

            if not hasattr(self, "_crash_print_count"):
                self._crash_print_count = 0
            if self._crash_print_count < 5:
                print(
                    f"[FEM Crash] step={self._timestep}, "
                    f"elements_before={self._previous_num_elements if self._previous_num_elements is not None else '?'}, "
                    f"type={type(e).__name__}: {str(e)[:300]}",
                    file=sys.stderr,
                    flush=True,
                )
                self._crash_print_count += 1
            elif self._crash_print_count == 5:
                print(
                    f"[FEM Crash] ... suppressing further crash prints ...",
                    file=sys.stderr,
                    flush=True,
                )
                self._crash_print_count += 1
            crash_penalty = -10000.0
            # 使用细化前的元素数，确保 reward 尺寸与 policy 输出的 values 一致
            # self._previous_num_elements 已在 _set_previous_step() 中保存
            num_agents = (
                self._previous_num_elements
                if self._previous_num_elements is not None
                else (self.num_agents if self.num_agents > 0 else 1)
            )
            self._reward = np.full(num_agents, crash_penalty, dtype=np.float32)
            self._cumulative_return = self._cumulative_return + np.sum(self._reward)
            # 确保 agent_mapping 与 reward/values 维度一致
            self._agent_mapping = np.arange(num_agents, dtype=np.int64)
            # _num_elements is a property, cannot be set directly

            # 返回当前观测 (如果可用) 或空图
            try:
                observation = self.last_observation
            except Exception:
                # 创建一个最小空图作为 fallback
                observation = Data(
                    x=torch.zeros(
                        (num_agents, self.num_node_features), dtype=torch.float32
                    ),
                    edge_index=torch.zeros((2, 0), dtype=torch.long),
                    edge_attr=torch.zeros(
                        (0, self.num_edge_features), dtype=torch.float32
                    ),
                )

            info = {
                "is_truncated": False,
                "return": float(np.sum(self._reward)),
                "weighted_remaining_error": float("inf"),
                "num_elements": self._num_elements
                if self.fem_problem is not None
                else 0,
                "num_agents": num_agents,
                "fem_crash": True,
                "crash_reason": str(e)[:200],  # 截断以防日志过长
            }
            self._timestep = self._max_timesteps  # 强制终止
            return observation, self._reward, True, info

    def inference_step(
        self, action: np.ndarray
    ) -> Tuple[Data, float, bool, Dict[str, Any]]:
        """
        Performs a step of the Mesh Refinement task *without* calculating the reward or difference to the fine-grained
        reference. This is used for inference

        Args:
            action: the action the agents will take in this step. Has shape (num_agents, action_dimension)
            Given as an array of shape (num_agents, action_dimension)

        Returns: A 4-tuple (observations, reward, done, info), where
            * observations is a graph of the agents and their positions, in this case of the refined mesh
            * reward is a single scalar shared between all agents, i.e., per **graph**
            * done is a boolean flag that says whether the current rollout is finished or not
            * info is a dictionary containing additional information
        """
        assert not self.is_terminal, (
            f"Tried to perform a step on a terminated environment. Currently on step "
            f"{self._timestep:} of {self._max_timesteps:} "
            f"with {self._num_elements}/{self._maximum_elements} elements."
        )

        self._timestep += 1
        self._agent_mapping = self._refine_mesh(action=action)
        # solve equation
        self.fem_problem.calculate_solution()
        observation = self.last_observation
        done = self.is_terminal
        info = {}
        return observation, self._reward, done, info

    def _set_previous_step(self):
        """
        Sets variables for the previous timestep. These are used for the reward function, as well as for different
        kinds of plots and metrics
        """
        self._previous_num_elements = self._num_elements
        self._previous_agent_mapping = self._agent_mapping
        self._previous_element_volumes = self.element_volumes
        self._previous_eta_indicator = self._eta_indicator
        self._previous_eta_components = self._eta_components_raw
        self._previous_rho_components = self._rho_components
        self._previous_solution_l2_norm = self._compute_solution_l2_norm()

    def _compute_solution_l2_norm(self) -> float:
        """Approximate ||u_h||_{L2(Ω)} via element centroids: sqrt(Σ_K |ū_K|² · area_K)."""
        u_scat = self.fem_problem.nodal_solution  # complex (n_vertices,)
        elem_idx = self._element_indices           # (n_elements, 3)
        vols = self.element_volumes                # (n_elements,)
        u_elem = u_scat[elem_idx]                  # (n_elements, 3)
        u_elem_mean = np.mean(u_elem, axis=1)      # (n_elements,) complex mean
        u_mag = np.abs(u_elem_mean)
        return float(np.sqrt(np.sum(u_mag ** 2 * vols)))

    def _refine_mesh(self, action: np.array) -> np.array:
        """
        Refines fem_problem.mesh by splitting all faces/elements for which the average of agent activation surpasses a
        threshold.
        If this refinement exceeds the maximum number of nodes allowed in the environment, we return a boolean flag
        that indicates so and stops the environment

        Optionally smoothens the newly created mesh as a post-processing step

        Args:
            action: An action/activation per element.
                - continuous_sizing_field: shape (num_agents, 1) or (num_agents,) → 目标网格面积
                - absolute/absolute_discrete: shape (num_agents,) or (num_agents, 1) → scalar threshold
        Returns: An array of mapped element indices

        """
        # 标量动作统一 flatten 到 1D
        action = action.flatten()

        elements_to_refine = self._get_elements_to_refine(action)

        # updates self.fem_problem.mesh
        element_mapping = self.fem_problem.refine_mesh(elements_to_refine)
        return element_mapping

    def _get_elements_to_refine(self, action: np.array) -> np.array:
        """
        Calculate which elements to refine based on the action, refinement strategy and the
        maximum number of elements allowed in the environment
        Args:
            action: An action/activation per agent, i.e., per element. 1D array of shape (num_agents,).
                - continuous_sizing_field: 每个 agent 输出 1 个标量 → Softplus → 期望最大单元面积
                - absolute/absolute_discrete: scalar threshold

        Returns: An array of ids corresponding to elements_to_refine

        """
        # select elements to refine based on the average actions of its surrounding agents/nodes

        if self._refinement_strategy == "continuous_sizing_field":
            # ================================================================
            #  Score-based 细化选择（由 actor 直接排序，物理预算约束）
            #
            #  Actor 输出标量 x_i: x_i < 0 → 希望细化; x_i > 0 → 不希望细化
            #  排序依据 score = -x_i，在预算和上限内选 top-k
            #
            #  物理预算 N_budget: Σ area_K / A_budget，其中
            #    A_budget = ½(λ_local/6)²，对应每局部波长方向 ~6 个尺度点
            #
            #  动作掩码 (Dörfler-P95): η_K < 0.05·η_P95 的单元移出候选池，
            #   P95 锚定物理误差尺度，免疫远场噪声稀释，强制预算投入误差主导区
            # ================================================================
            x = action.flatten()

            # ── 训练监控指标（在所有 early return 之前计算）──
            self._diag_neg_ratio = float(np.mean(x < 0.0))

            remaining = self._n_budget - self._num_elements
            max_parents_by_budget = max(0, remaining // 3)

            if max_parents_by_budget <= 0:
                self._diag_eligible_ratio = 0.0
                self._diag_selected_count = 0
                return np.array([], dtype=np.int64)

            # 动态计算每单元预算面积（仅用于 N_budget 全局资源上限）
            eps_r_elem = self._epsilon_r_elements
            k = self._wave_number
            lambda_local = 2.0 * np.pi / (k * np.sqrt(np.maximum(eps_r_elem, 1.0)))
            A_budget = 0.5 * (lambda_local / 6.0) ** 2

            # 纯数值安全底线：仅防止 scikit-fem 因浮点精度导致的退化/奇异。
            # 不再用 0.25*A_budget —— RL 应自主学会"多细才够"，
            # 而非被人为启发式 (12 点/波长) 限制。
            domain_area = float(np.prod(self.fem_problem.plot_boundary[2:] - self.fem_problem.plot_boundary[:2]))
            V_min_safeguard = 1e-10 * domain_area

            # Filter 1: numerical safeguard only — no physics heuristic
            area_eligible = np.where(self.element_volumes > V_min_safeguard)[0]

            # Filter 2: Dörfler-style action mask — exclude elements below 5% of η_P95
            # P95 anchors the threshold to physically meaningful error scale,
            # immune to far-field noise dilution (unlike median or mean).
            # η_K < 0.05·η_P95 → not worth the refinement budget.
            eta_current = self._eta_indicator
            eta_p95 = np.percentile(eta_current, 95)
            error_eligible = np.where(eta_current >= 0.05 * eta_p95)[0]

            eligible = np.intersect1d(area_eligible, error_eligible)

            self._diag_eligible_ratio = float(len(eligible)) / max(self._num_elements, 1)
            self._diag_masked_ratio = (
                1.0 - float(len(eligible)) / max(len(area_eligible), 1)
                if len(area_eligible) > 0 else 0.0
            )

            num = min(
                len(eligible),
                max(1, self._num_elements // 4),
                max_parents_by_budget,
            )

            if num <= 0:
                self._diag_selected_count = 0
                return np.array([], dtype=np.int64)

            # x 越小 ⇒ 优先级越高（纯排序，不设正负门槛）
            score = -x
            selected = eligible[np.argsort(score[eligible])[-num:]]
            self._diag_selected_count = len(selected)
            elements_to_refine = selected

        elif self._refinement_strategy in ["absolute", "absolute_discrete"]:
            elements_to_refine = np.argwhere(action > 0.0).flatten()
        else:
            raise ValueError(
                f"Unknown refinement strategy '{self._refinement_strategy}"
            )
        return elements_to_refine

    def render(self, mode: str = "human", render_intermediate_steps: bool = False, *args, **kwargs):
        if not (render_intermediate_steps or self.is_terminal):
            return [], {}
        remaining_error = self._get_remaining_error(return_dimensions=False)
        title = (
            f"Solution. Element Penalty: {self._element_penalty_lambda:.1e} "
            f"Reward: {np.sum(self._reward):.3f} Return: {np.sum(self._cumulative_return):.3f} "
            f"Agents: {self.num_agents} Remaining Error: {remaining_error:.3f}"
        )
        traces, layout = get_plotly_mesh_traces_and_layout(
            mesh=self.mesh, scalars=np.real(self.scalar_solution),
            mesh_dimension=2, title=title, boundary=self.fem_problem.plot_boundary,
        )
        _fp = self.fem_problem.fem_problem
        cx = getattr(_fp, "_cx", 0.5)
        cy = getattr(_fp, "_cy", 0.5)
        r = getattr(_fp, "_radius", 0.2)
        traces.append(go.Scatter3d(
            x=cx + r * np.cos(np.linspace(0, 2 * np.pi, 128)),
            y=cy + r * np.sin(np.linspace(0, 2 * np.pi, 128)),
            z=np.zeros(128), mode="lines",
            line=dict(color="cyan", width=2, dash="dash"),
            name="Dielectric", showlegend=True,
        ))
        return traces, {"layout": layout}

    def _get_remaining_error(
        self, return_dimensions: bool = False
    ) -> Union[np.array, Tuple]:
        """
        Get the remaining error by aggregating over all elements and taking the convex sum of all solution dimensions
        """
        remaining_error_per_dimension = np.sqrt(
            np.sum(self.error_per_element**2, axis=0)
        )

        # Collapse per-element/per-dim initial error to scalar if needed
        norm = np.atleast_1d(np.asarray(self.initial_approximation_error, dtype=np.float64))
        if remaining_error_per_dimension.ndim < norm.ndim or (
            remaining_error_per_dimension.ndim == norm.ndim
            and remaining_error_per_dimension.shape != norm.shape
        ):
            norm = np.sqrt(np.sum(norm**2))

        remaining_error_per_dimension = (
            remaining_error_per_dimension / norm
        )  # normalize
        remaining_error = self.project_to_scalar(remaining_error_per_dimension)
        # Ensure scalar output (defensive against (1,) or (Ne,) arrays from 1D PDEs)
        remaining_error = float(np.asarray(remaining_error).ravel()[0])

        if return_dimensions:
            return remaining_error, remaining_error_per_dimension
        else:
            return remaining_error

    def _compute_total_error(self, error_per_element: np.ndarray = None) -> float:
        """
        Compute a scalar total error from a per-element error array.
        Uses the same aggregation (sum or max) as _get_remaining_error,
        but without normalization by the initial approximation error.

        Args:
            error_per_element: Per-element error array of shape (num_elements, solution_dimension).
                If None, uses the current error_per_element.

        Returns: A scalar total error.
        """
        if error_per_element is None:
            error_per_element = self.error_per_element
        error_per_dim = np.sqrt(np.sum(error_per_element**2, axis=0))
        return float(self.project_to_scalar(error_per_dim))

    @property
    def last_observation(self) -> Data:
        """
        Retrieve an observation graph for the current state of the environment.

        We use an additional self.last_observation wrapper to make sure that classes that inherit from this
        one have access to node and edge features outside the Data() structure
        Returns: A Data() object of the graph that describes the current state of this environment

        """
        graph_dict = {}
        graph_dict = graph_dict | self._get_graph_nodes()
        graph_dict = graph_dict | self._get_graph_edges()

        observation_graph = Data(**graph_dict)

        return observation_graph

    def _get_graph_nodes(self) -> Dict[str, Dict[str, torch.Tensor]]:
        """
        Returns a dictionary of node features that are used to describe the current state of this environment.

        Returns: A dictionary of node features. This dictionary has the format
        {"x": element_features}
        where element and node features depend on the context, but include things like the evaluation of the target
        function, the degree of the node, etc.
        """
        # Builds feature matrix of shape (#elements, #features)
        # by iterating over the functions in self._element_feature_functions.
        element_features = np.array(
            [fn() for key, fn in self._element_feature_functions.items()]
        ).T
        element_features = save_concatenate(
            [element_features, self.fem_problem.element_features()], axis=1
        )
        element_features = torch.tensor(element_features, dtype=torch.float32)
        node_dict = {"x": element_features}
        return node_dict

    def _get_graph_edges(
        self,
    ) -> Dict[Union[str, Tuple[str, str, str]], Dict[str, torch.Tensor]]:
        """
        Returns a dictionary of edge features that are used to describe the current state of this environment.
        Note that we always use symmetric graphs and self edges.

        Returns: A dictionary of edge features. This dictionary has the format
        {
        "edge_index": indices,
        "edge_attr": features
        }
        """
        edge_index, edge_attr = self._element2element_features(self._num_edge_features)
        edge_dict = {"edge_index": edge_index, "edge_attr": edge_attr}
        return edge_dict

    def _element2element_features(self, num_edge_features: int):
        # concatenate incoming, outgoing and self edges of each node to get an undirected graph
        src_nodes = np.concatenate(
            (
                self._element_neighbors[0],
                self._element_neighbors[1],
                np.arange(self._num_elements),
            ),
            axis=0,
        )
        dest_nodes = np.concatenate(
            (
                self._element_neighbors[1],
                self._element_neighbors[0],
                np.arange(self._num_elements),
            ),
            axis=0,
        )
        num_edges = self._element_neighbors.shape[1] * 2 + self._num_elements
        edge_features = np.empty(shape=(num_edges, num_edge_features))
        edge_feature_position = 0
        if "euclidean_distance" in self._edge_features:
            euclidean_distances = np.linalg.norm(
                self._element_midpoints[dest_nodes]
                - self._element_midpoints[src_nodes],
                axis=1,
            )
            lam = 2.0 * np.pi / self._wave_number
            edge_features[:, edge_feature_position] = euclidean_distances / lam
            edge_feature_position += 1
        edge_index = torch.tensor(np.vstack((src_nodes, dest_nodes))).long()
        edge_attr = torch.tensor(edge_features, dtype=torch.float32)
        return edge_index, edge_attr

    def _get_reward_dict(self) -> Dict[str, np.float32]:
        """
        Calculate the reward for the current timestep depending on the environment states and the action
        the agents took.
        Args:

        Returns:
            Dictionary that must contain "reward" as well as partial reward data

        """
        reward, reward_dict = self._get_reward_by_type()

        self._reward = reward
        self._cumulative_return = self._cumulative_return + np.sum(self._reward)

        return reward_dict

    def _get_metric_dict(self) -> Dict[str, Any]:
        remaining_error, remaining_error_per_dimension = self._get_remaining_error(
            return_dimensions=True
        )

        metric_dict = {
            "weighted_remaining_error": remaining_error,
            "error_times_agents": remaining_error * self.num_agents,
            "delta_elements": self._num_elements - self._previous_num_elements,
            "avg_total_refinements": np.log(
                self._num_elements / self._initial_num_elements
            )
            / np.log(4),
            "avg_step_refinements": np.log(
                self._num_elements / self._previous_num_elements
            )
            / np.log(4),
            "num_elements": self._num_elements,
            "num_agents": self.num_agents,
            "reached_element_limits": self.reached_element_limits,
            "refinement_std": self._refinements_per_element.std(),
        }

        for error_metric, element_errors in self.error_estimation_dict.items():
            if element_errors.ndim >= 1 and element_errors.shape[0] == self._num_elements:
                error_per_dimension = np.max(element_errors, axis=0)
            else:
                error_per_dimension = element_errors
            error_per_dimension = (
                error_per_dimension / self._initial_approximation_errors[error_metric]
            )
            remaining_error = self.project_to_scalar(error_per_dimension)
            metric_dict[f"{error_metric}_error"] = remaining_error
        return metric_dict

    def _get_action_dict(self, action: np.ndarray) -> Dict[str, Any]:
        """
        Returns a dictionary of information about the action that was taken in the current timestep
        Args:
            action: The action that was taken in the current timestep

        Returns: A dictionary of information about the action that was taken in the current timestep

        """
        action_dict = {}
        if self._refinement_strategy in ["absolute", "absolute_discrete"]:
            action_dict["action_mean"] = np.mean(action)
            action_dict["action_std"] = np.std(action)
        return action_dict

    def _get_reward_by_type(self) -> Tuple[np.array, Dict]:
        """
        Potential-based reward shaping on η indicator.

        spatial_max — Per-agent reward (parent i → children C(i)):
          r_local_i = log(η_old_i + ε_dynamic) − log(max_{j∈C(i)} η_new_j + ε_dynamic)
                      − λ·(|C(i)| − 1)

        spatial — Per-agent reward (parent i → children C(i)):
          r_local_i = log(η_old_i + ε_dynamic) − log(√(Σ_{j∈C(i)} η_new_j²) + ε_dynamic)
                      − λ·(|C(i)| − 1)

        ε_dynamic = max(0.01 × η_P95, 1e-6) — anchored to P95 of residual,
        immune to far-field dilution; prevents reward hacking on near-zero-η elements.

        Potential function:  Φ(s) = −log(E_global)
          E_global = √(Σ_K η_K²) / ||u_h||_{L2(Ω)}     (dimensionless)
        Shaped reward:       r_i = r_local_i + α · (log E_old − log E_new)
        """

        reward_dict = {}
        # Dynamic epsilon anchored to P95 of η — immune to far-field dilution
        # that plagues mean-based approaches. P95 is driven by physically
        # meaningful error in the dielectric, not background noise.
        # ε_dynamic = max(0.01 × η_P95, 1e-6)
        eta_current_raw = self._eta_indicator
        eta_p95 = float(np.percentile(eta_current_raw, 95))
        eps = max(0.01 * eta_p95, 1e-6)

        old_eta = self._previous_eta_indicator + eps
        new_eta = eta_current_raw + eps

        if self._reward_type == "spatial_max":
            from torch_scatter import scatter_max

            agent_mapping = torch.tensor(self.agent_mapping)
            child_eta = torch.tensor(new_eta)
            max_child_eta, _ = scatter_max(
                src=child_eta,
                index=agent_mapping,
                dim=0,
                dim_size=old_eta.shape[0],
            )
            max_child_eta = max_child_eta.numpy()
            reward_per_agent_and_dim = np.log(old_eta) - np.log(max_child_eta)

        elif self._reward_type == "spatial":
            from torch_scatter import scatter_add

            agent_mapping = torch.tensor(self.agent_mapping)
            # L₂ aggregation: √(Σ η_child²) — never punishes refinement
            child_eta = torch.tensor(new_eta)
            sum_sq_child_eta = scatter_add(
                src=child_eta * child_eta,
                index=agent_mapping,
                dim=0,
                dim_size=old_eta.shape[0],
            )
            l2_child_eta = np.sqrt(sum_sq_child_eta.numpy()) + eps
            reward_per_agent_and_dim = np.log(old_eta) - np.log(l2_child_eta)

        else:
            raise ValueError(f"Unknown reward type {self._reward_type}")

        reward_per_agent = self.project_to_scalar(reward_per_agent_and_dim)

        # apply action/element penalty
        unique_old, counts = np.unique(self.agent_mapping, return_counts=True)
        element_penalty = np.zeros(len(reward_per_agent), dtype=reward_per_agent.dtype)
        element_penalty[unique_old] = self._element_penalty_lambda * (counts - 1)
        element_limit_penalty = (
            (self._element_limit_penalty / self._previous_num_elements)
            if self.reached_element_limits
            else 0
        )
        reward_per_agent = (
            reward_per_agent - element_penalty - element_limit_penalty
        )

        # ── Potential-based shaping: only refined parents get the global bonus ──
        l2_old = self._previous_solution_l2_norm
        l2_new = self._compute_solution_l2_norm()
        eta_l2_old = float(np.sqrt(np.sum(old_eta ** 2)))
        eta_l2_new = float(np.sqrt(np.sum(new_eta ** 2)))
        eps_l2 = 1e-12
        E_old = eta_l2_old / max(l2_old, eps_l2)
        E_new = eta_l2_new / max(l2_new, eps_l2)
        global_bonus = self._global_reward_alpha * float(np.log(E_old + eps_l2) - np.log(E_new + eps_l2))
        refined_parents = unique_old[counts > 1]
        reward_per_agent[refined_parents] += global_bonus

        self._reward_per_agent = reward_per_agent
        self._cumulative_reward_per_agent = (
            self._cumulative_reward_per_agent[self._previous_agent_mapping]
            + reward_per_agent
        )
        reward = reward_per_agent

        reward_dict["reward"] = reward
        reward_dict["potential_bonus"] = global_bonus
        reward_dict["penalty"] = -reward
        reward_dict["element_limit_penalty"] = element_limit_penalty
        reward_dict["element_penalty"] = element_penalty
        reward_dict["element_penalty_lambda"] = self._element_penalty_lambda
        return reward, reward_dict

    @property
    def mesh(self) -> Mesh:
        """
        Returns the current mesh.
        """
        return self.fem_problem.mesh

    @property
    def agent_node_type(self) -> str:
        return "element"

    @property
    def _vertex_positions(self) -> np.array:
        """
        Returns the positions of all vertices/nodes of the mesh.
        Returns: np.array of shape (num_vertices, 2)
        """
        return self.fem_problem.vertex_positions

    @property
    def _element_indices(self) -> np.array:
        return self.fem_problem.element_indices

    @property
    def _element_midpoints(self) -> np.array:
        """
        Returns the midpoints of all elements/faces.
        Returns: np.array of shape (num_elements, 2)

        """
        return self.fem_problem.element_midpoints

    @property
    def _mesh_edges(self) -> np.array:
        """
        Returns: the edges of all vertices/nodes of the mesh. Shape (2, num_edges)
        """
        return self.fem_problem.mesh_edges

    @property
    def _element_neighbors(self) -> np.array:
        """
        Find neighbors of each element. Shape (2, num_neighbors)
        Returns:

        """
        # f2t are element/face neighborhoods, which are set to -1 for boundaries
        return self.fem_problem.element_neighbors

    @property
    def _num_elements(self) -> int:
        return len(self._element_indices)

    @property
    def _num_vertices(self) -> int:
        return len(self._vertex_positions)

    @property
    def element_volumes(self) -> np.array:
        return get_triangle_areas_from_indices(
            positions=self._vertex_positions, triangle_indices=self._element_indices
        )

    @property
    def num_node_features(self) -> int:
        return self._num_node_features

    @property
    def num_edge_features(self) -> int:
        return self._num_edge_features

    @property
    def action_dimension(self) -> int:
        """
        Returns: The dimensionality of the action space.

        - continuous_sizing_field: 1D continuous output → 目标网格面积 (Softplus 激活)
        - absolute_discrete: 2 discrete actions (refine / don't refine)
        - others: single continuous scalar
        """
        if self._refinement_strategy == "continuous_sizing_field":
            return 1  # 1D 连续标量 → Softplus → 目标面积 S_i
        elif self._refinement_strategy == "absolute_discrete":
            return 2
        else:  # single continuous value
            return 1

    @property
    def num_agents(self) -> int:
        if self.fem_problem is not None and self.fem_problem.mesh is not None:
            return self._num_elements
        else:
            return 1  # placeholder

    @property
    def _action_space(self) -> gym.Space:
        """

        Returns: The **current** action space of the environment. Bound to change, since the number of agents
        changes

        """
        if self._refinement_strategy in ["absolute_discrete", "argmax", "single_agent"]:
            return gym.spaces.MultiDiscrete([self.action_dimension] * self.num_agents)
        elif self._refinement_strategy == "continuous_sizing_field":
            # 连续 1D 输出: 每个 agent 输出 1 个标量 → Softplus → 目标网格面积
            # 无界连续空间，PPO Gaussian policy 负责探索
            return gym.spaces.Box(
                low=-1e5,
                high=1e5,
                shape=(self.num_agents, self.action_dimension),
                dtype=np.float32,
            )
        else:
            return gym.spaces.Box(
                low=-1e5,
                high=1e5,
                shape=(
                    self.num_agents,
                    self.action_dimension,
                ),
                dtype=np.float32,
            )

    @property
    def agent_mapping(self) -> np.array:
        assert self._agent_mapping is not None, "Element mapping not initialized"
        return self._agent_mapping

    @property
    def previous_agent_mapping(self) -> np.array:
        assert self._previous_agent_mapping is not None, (
            "Previous element mapping not initialized"
        )
        return self._previous_agent_mapping

    @property
    def reached_element_limits(self) -> bool:
        """
        True if the number of elements/faces in the mesh is above the maximum allowed value.
        Returns:

        """
        return self._num_elements > self._maximum_elements

    @property
    def is_truncated(self) -> bool:
        return self._timestep >= self._max_timesteps

    @property
    def is_terminal(self) -> bool:
        # Agent selected nothing to refine — budget exhausted or
        # Doerfler mask filtered everything.  Episode converged naturally.
        # -1 = not yet evaluated (reset state), 0 = nothing selected this step.
        sc = getattr(self, "_diag_selected_count", -1)
        if sc == 0:
            return True
        return self.reached_element_limits or self.is_truncated

    @property
    def solution(self) -> np.array:
        """
        Returns: solution vector per *vertex* of the mesh.
            An array (num_vertices, solution_dimension),
            where every entry corresponds to the solution of the parameterized fem_problem
            equation at the position of the respective node/vertex.

        """
        return self.fem_problem.nodal_solution

    def project_to_scalar(self, values: np.array) -> np.array:
        """
        Projects a value per node or graph and solution dimension to a scalar value per node.
        Args:
            values: A vector of shape ([num_vertices/nodes,] solution_dimension)

        Returns: A scalar value per vertex
        """
        return self.fem_problem.project_to_scalar(values)

    @property
    def scalar_solution(self):
        return self.project_to_scalar(self.solution)

    @property
    def error_per_element(self) -> np.array:
        """
        Returns: error per element of the mesh. np.array of shape (num_elements, solution_dimension)

        """
        return self._error_estimation_dict.get("indicator")

    @property
    def initial_approximation_error(self) -> np.array:
        """
        Returns: error per element of the mesh. np.array of shape (num_elements, solution_dimension)

        """
        return self._initial_approximation_errors.get("indicator")

    @property
    def error_estimation_dict(self) -> Dict[str, np.array]:
        """
        Returns a dictionary of all error estimation methods and their respective errors.
        These errors may be per element/face, or per integration point, depending on the metric.
        Returns:

        """
        return self._error_estimation_dict

    @property
    def _refinements_per_element(self) -> np.array:
        return self.fem_problem.refinements_per_element

    @property
    def _solution_std_per_element(self) -> np.array:
        """
        Computes the standard deviation of the solution per element.
        Returns: np.array of shape (num_elements, solution_dimension)

        Note: 此属性仅用于 backward compatibility；
        新代码使用 _residual_components 替代。
        """
        return get_aggregation_per_element(
            self.solution, self._element_indices, aggregation_function_str="std"
        )

    # =========================================================================
    #  PDE 物理残差特征 (替代 solution_std)
    # =========================================================================

    @property
    def _residual_components(self) -> Dict[str, np.ndarray]:
        """逐单元的三项 PDE 残差 + 边界标记。"""
        from .helmholtz import _compute_residual_components

        fp = self.fem_problem.fem_problem
        k = getattr(fp, "_k", 10.0)
        u_scat = self.fem_problem.nodal_solution
        eps_r = self._epsilon_r_elements
        return _compute_residual_components(
            self.fem_problem.mesh, u_scat, k=k, eps_r=eps_r
        )

    @property
    def _k_local_sqrt_vol(self) -> np.ndarray:
        """每个单元的 k_local × sqrt(volume)。"""
        k = self._wave_number
        eps_r = self._epsilon_r_elements
        k_local = k * np.sqrt(np.maximum(eps_r, 1.0))
        return (k_local * np.sqrt(self.element_volumes)).astype(np.float32)

    @property
    def _volume_normalized(self) -> np.ndarray:
        """无量纲单元面积: volume / lambda^2。"""
        lam = 2.0 * np.pi / self._wave_number
        return (self.element_volumes / (lam * lam)).astype(np.float32)

    @property
    def _eta_indicator(self) -> np.ndarray:
        """
        标准 FEM 残差误差指示器，用于 reward 计算。

        η_i = √(R_int_i² + J_grad_i² + R_sbc_i²)

        其中:
          R_int_i = h_K · √V_i · |k²ε_r u + k²(ε_r-1)u_inc|
          J_grad_i = √(½ Σ_{e∈∂K_i} h_e² · |[[∇u·n]]|²)
          R_sbc_i = √h_bnd · |∂u/∂n - i·k_local·u|

        与 _compute_residual_indicator 的公式完全一致。

        Returns: shape (num_elements,) float64
        """
        from .helmholtz import _compute_residual_components

        fp = self.fem_problem.fem_problem
        k = getattr(fp, "_k", 10.0)
        u_scat = self.fem_problem.nodal_solution
        comps = _compute_residual_components(
            self.fem_problem.mesh, u_scat, k=k,
            eps_r=self._epsilon_r_elements, apply_log=False,
        )
        self._cached_eta_components_raw = comps
        return np.sqrt(
            comps["internal_residual"] ** 2
            + comps["gradient_jump"] ** 2
            + comps["sbc_residual"] ** 2
        )

    @property
    def _eta_components_raw(self) -> Dict[str, np.ndarray]:
        """返回逐单元的三项原始残差分量（apply_log=False），由 _eta_indicator 缓存。"""
        if not hasattr(self, "_cached_eta_components_raw") or self._cached_eta_components_raw is None:
            _ = self._eta_indicator  # triggers caching
        return self._cached_eta_components_raw

    @property
    def _rho_components(self) -> Dict[str, np.ndarray]:
        """返回逐单元的残差密度三分量（不含 h-缩放），用于 reward 计算。

        Returns:
            rho_int:  |k²·ε_r·u + k²·(ε_r-1)·u_inc|
            rho_jump: √(mean |[[∇u·n]]|²) per element
            rho_sbc:  √(mean |∂u/∂n - i·k·u|²) per element
        """
        from .helmholtz import _compute_residual_density

        fp = self.fem_problem.fem_problem
        k = getattr(fp, "_k", 10.0)
        u_scat = self.fem_problem.nodal_solution
        return _compute_residual_density(
            self.fem_problem.mesh, u_scat, k=k,
            eps_r=self._epsilon_r_elements,
        )

    # =========================================================================
    #  SBC 状态空间辅助属性：介电常数 + 复数场均值
    # =========================================================================

    @property
    def _wave_number(self) -> float:
        """Helmholtz 波数 k，从当前 FEM 问题实例读取（支持随机采样）。"""
        fp = self.fem_problem.fem_problem
        return getattr(fp, '_k', 10.0)

    @property
    def _epsilon_r_elements(self) -> np.ndarray:
        """
        每个单元的相对介电常数 εr。

        从 FEM 问题实例读取介质几何参数，按单元中点判断是否在介质内。

        Returns: shape (num_elements,) float64 array
        """
        fp = self.fem_problem.fem_problem
        cx = getattr(fp, "_cx", 0.5)
        cy = getattr(fp, "_cy", 0.5)
        radius = getattr(fp, "_radius", 0.2)
        eps_r = getattr(fp, "_eps_r", 2.0)
        midpoints = self._element_midpoints
        x_mid, y_mid = midpoints[:, 0], midpoints[:, 1]
        in_cylinder = (x_mid - cx) ** 2 + (y_mid - cy) ** 2 <= radius**2
        return np.where(in_cylinder, eps_r, 1.0)

    @property
    def _dist_to_interface(self) -> np.ndarray:
        """每个单元中点到介质圆柱边界的带符号距离（内部为负，外部为正）。

        用真空波长 lambda = 2*pi/k 做无量纲归一化，再经 sign(d)·ln(1+|d|) 压缩。
        ln 压缩保留近场分辨力（小 |d| 时近似线性），远场自然对数压缩，
        与残差特征的 log₁₀ 压缩风格一致。无硬截断，处处可导。
        """
        fp = self.fem_problem.fem_problem
        cx = getattr(fp, "_cx", 0.5)
        cy = getattr(fp, "_cy", 0.5)
        radius = getattr(fp, "_radius", 0.2)
        midpoints = self._element_midpoints
        dist = np.sqrt((midpoints[:, 0] - cx) ** 2 + (midpoints[:, 1] - cy) ** 2)
        lam = 2.0 * np.pi / self._wave_number
        d = (dist - radius) / lam
        return (np.sign(d) * np.log1p(np.abs(d))).astype(np.float32)

    @property
    def _eps_r_global(self) -> float:
        """散射体材料的相对介电常数（全局标量）。"""
        fp = self.fem_problem.fem_problem
        return getattr(fp, "_eps_r", 2.0)

    @property
    def _complex_solution_mean(self) -> np.ndarray:
        """
        每个单元内复数 FEM 解的均值 (complex128)。

        SBC 边界条件下解为复数值；内部残差和边界残差均基于复数场。
        使用 P1 节点值的三点平均作为单元代表值。

        Returns: shape (num_elements,) complex128 array
        """
        return get_aggregation_per_element(
            self.solution,
            self._element_indices,
            aggregation_function_str="mean",
        )

    @property
    def sample_penalty(self) -> bool:
        return self._sample_penalty

    @property
    def refinement_strategy(self) -> str:
        return self._refinement_strategy

    @property
    def has_homogeneous_graph(self) -> bool:
        return not self._include_vertices

    @property
    def mesh_dimension(self) -> int:
        return 2

    def set_element_penalty_lambda(
        self, position_or_value: float, from_position: bool = True
    ):
        """
        Sets the element penalty lambda from the provided position.
        Args:
            position_or_value: A float between 0 and 1 that determines the element penalty lambda if from_position.
                Otherwise, the value of the element penalty lambda.
            from_position: If True, the element penalty lambda is taken log-uniformly from the provided position,
                regardless of how the value is usually sampled.


        Returns: None

        Note: Sets self._element_penalty_lambda

        """
        element_penalty_config = self._environment_config.get("element_penalty")

        if element_penalty_config.get("sample_penalty"):
            if from_position:
                # sample element penalty loguniformly for comparison between methods
                log_min = np.log(element_penalty_config.get("min"))
                log_max = np.log(element_penalty_config.get("max"))
                self._element_penalty_lambda = np.exp(
                    position_or_value * log_min + (1 - position_or_value) * log_max
                )
            else:  # fixed element penalty
                self._element_penalty_lambda = position_or_value
        else:
            # element penalty is fixed
            self._element_penalty_lambda = element_penalty_config.get("value")

    ####################
    # additional plots #
    ####################

    def _plot_value_per_element(
        self,
        value_per_element: np.array,
        title: str,
        normalize_by_element_volume: bool = False,
        mesh: Optional[Mesh] = None,
    ) -> go.Figure:
        """
        only return traces if asked or at the last step to avoid overlay of multiple steps
        Args:
            value_per_element: A numpy array of shape (num_elements,).
            title: The title of the plot.
            normalize_by_element_volume: If True, the values are normalized by the element volume as value /= element_volume.
            mesh: The mesh to plot the values on. If None, the mesh of the current state is used.
        Returns: A plotly figure with an outline of the mesh and value per element in the element midpoints.

        """
        if mesh is None:
            assert len(value_per_element) == self.num_agents, (
                f"Need to provide a value per agent, given "
                f"'{value_per_element.shape}' and '{self.num_agents}'"
            )
            mesh = self.fem_problem.mesh
            mesh_dimension = 2
        else:
            mesh_dimension = mesh.dim()

        if normalize_by_element_volume:
            value_per_element = value_per_element / self.element_volumes

        boundary = self.fem_problem.plot_boundary
        traces, layout = get_plotly_mesh_traces_and_layout(
            mesh=mesh,
            scalars=value_per_element,
            mesh_dimension=mesh_dimension,
            title=title,
            boundary=boundary,
        )

        value_per_element_plot = go.Figure(data=traces, layout=layout)
        return value_per_element_plot

    def _plot_error_per_element(
        self, normalize_by_element_volume: bool = True
    ) -> go.Figure:
        weighted_remaining_error = self._get_remaining_error(return_dimensions=False)
        return self._plot_value_per_element(
            value_per_element=self.project_to_scalar(self.error_per_element),
            normalize_by_element_volume=normalize_by_element_volume,
            title=f"Element Errors. Remaining total error: {weighted_remaining_error:.4f}",
        )

    def additional_plots(
        self, iteration: int, policy_step_function: Optional[callable] = None
    ) -> Dict[str, go.Figure]:
        """
        Function that takes an algorithm iteration as input and returns a number of additional plots about the
        current environment as output. Some plots may be always selected, some only on e.g., iteration 0.
        Args:
            iteration: The current iteration of the algorithm.
            policy_step_function: (Optional)
                A function that takes a graph as input and returns the action(s) and (q)-value(s)
                for each agent.

        """
        _, remaining_error_per_solution_dimension = self._get_remaining_error(
            return_dimensions=True
        )
        additional_plots = {
            "refinements_per_element": self._plot_value_per_element(
                value_per_element=self._refinements_per_element,
                title="Refinements per element",
            ),
            "scalar_solution_std_per_element": self._plot_value_per_element(
                value_per_element=self.project_to_scalar(
                    self._solution_std_per_element
                ),
                title=f"Element Std of Solution Norm",
            ),
            "scalar_solution_error_per_element": self._plot_error_per_element(
                normalize_by_element_volume=False
            ),
        }

        if policy_step_function is not None:
            from .utils import detach

            actions, values = policy_step_function(observation=self.last_observation)
            if len(actions) == self._num_elements:
                additional_plots["final_actor_evaluation"] = (
                    self._plot_value_per_element(
                        detach(actions),
                        title=f"Action per Agent at step {self._timestep}",
                    )
                )
            if len(values) == self._num_elements:
                additional_plots["final_critic_evaluation"] = (
                    self._plot_value_per_element(
                        detach(values),
                        title=f"Critic Evaluation at step {self._timestep}",
                    )
                )
        if self._reward_type in ["spatial", "spatial_max", "spatial_volume"]:
            additional_plots["cumulative_reward"] = self._plot_value_per_element(
                value_per_element=self._cumulative_reward_per_agent,
                title="Cumulative Reward",
                mesh=self.fem_problem.previous_mesh,
            )
            additional_plots["reward_per_agent"] = self._plot_value_per_element(
                value_per_element=self._reward_per_agent,
                title="Final Reward",
                mesh=self.fem_problem.previous_mesh,
            )
        additional_plots |= self.fem_problem.additional_plots()
        return additional_plots

    def __deepcopy__(self, memo):
        """
        Overwrite deepcopy to reinitialize stateless (lambda-) functions
        it is sufficient to call the register functions,
        as only new objects for the stateless lambda functions have to be created
        Args:
            memo:

        Returns:

        """
        from copy import deepcopy

        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        for k, v in self.__dict__.items():
            setattr(result, k, deepcopy(v, memo))

        setattr(
            result, "_element_feature_functions", result._register_element_features()
        )
        return result
