import copy
import os
from typing import Any, Dict, List, Optional, Union

import numpy as np
from skfem import Basis, Mesh

from .fem_util import get_element_midpoints
from .helmholtz import HelmholtzProblem, create_helmholtz_problem
from .utils import IndexSampler


class FEMProblemWrapper:
    """Wraps a HelmholtzProblem, managing mesh, solution cache, and refinement history."""

    def __init__(
        self,
        *,
        fem_config: Dict[Union[str, int], Any],
        fem_problem: HelmholtzProblem,
        pde_features: Dict[str, List[str]],
    ):
        self._fem_config = fem_config
        self.fem_problem = fem_problem
        self._pde_element_feature_names = pde_features["element_features"]
        self._mesh: Optional[Mesh] = None
        self._previous_mesh: Optional[Mesh] = None
        self._solution: Optional[np.ndarray] = None
        self._nodal_solution: Optional[np.ndarray] = None
        self._refinements_per_element: Optional[np.ndarray] = None
        self._plot_boundary = np.array(fem_config.get("domain", {}).get("boundary", [0, 0, 1, 1]))

    def reset(self):
        self._mesh = self.fem_problem.initial_mesh
        self._previous_mesh = copy.deepcopy(self._mesh)
        self._refinements_per_element = np.zeros(self.num_elements, dtype=np.int32)

    def calculate_solution_and_get_error(self) -> Dict[str, np.ndarray]:
        self.calculate_solution()
        return self.get_error_estimate_per_element()

    def calculate_solution(self) -> None:
        self._solution = self.fem_problem.calculate_solution(basis=self._basis, cache=True)
        self._nodal_solution = self._solution

    def get_error_estimate_per_element(self) -> Dict[str, np.ndarray]:
        return self.fem_problem.get_error_estimate_per_element(
            basis=self._basis, solution=self._solution
        )

    def refine_mesh(self, elements_to_refine: np.ndarray) -> np.ndarray:
        if len(elements_to_refine) > 0:
            refined_mesh = self._mesh.refined(elements_to_refine)
            new_midpoints = refined_mesh.p[:, refined_mesh.t].mean(axis=1)
            element_finder = self._mesh.element_finder()
            corresponding_elements = element_finder(*new_midpoints)
            element_indices, inverse_indices, counts = np.unique(
                corresponding_elements, return_counts=True, return_inverse=True
            )
            self._refinements_per_element[element_indices] += counts - 1
            self._refinements_per_element = self._refinements_per_element[inverse_indices]
        else:
            refined_mesh = self._mesh
            inverse_indices = np.arange(self._mesh.t.shape[1]).astype(np.int64)

        self.mesh = refined_mesh
        return inverse_indices

    # ---- PDE 相关的单元特征（source_term 等）----
    def element_features(self) -> np.ndarray:
        return self.fem_problem.element_features(
            mesh=self._mesh, element_feature_names=self._pde_element_feature_names
        )

    # ---- 将多分量值归约为标量（Helmholtz 取实部）----
    def project_to_scalar(self, values: np.ndarray) -> np.ndarray:
        return self.fem_problem.project_to_scalar(values=values)

    # ---- 当前 FEM 网格 ----
    @property
    def mesh(self) -> Optional[Mesh]:
        return self._mesh

    @mesh.setter
    def mesh(self, mesh: Mesh) -> None:
        self._previous_mesh = copy.deepcopy(self._mesh)
        self._mesh = mesh

    # ---- P1 线性基函数 ----
    @property
    def _basis(self) -> Basis:
        return self.fem_problem.mesh_to_basis(self._mesh)

    # ---- 细化前的网格（奖励计算中回溯用）----
    @property
    def previous_mesh(self) -> Mesh:
        return self._previous_mesh

    # ---- 当前网格单元总数 ----
    @property
    def num_elements(self) -> int:
        return self._mesh.t.shape[1]

    # ---- 每个单元被细化的次数 ----
    @property
    def refinements_per_element(self) -> np.ndarray:
        return self._refinements_per_element

    # ---- 顶点上的 FEM 解 ----
    @property
    def nodal_solution(self) -> np.ndarray:
        assert self._nodal_solution is not None, "Solution not computed yet"
        return self._nodal_solution

    # ---- 单元中点坐标 (num_elements, 2) ----
    @property
    def element_midpoints(self) -> np.ndarray:
        return get_element_midpoints(self._mesh)

    # ---- 单元顶点索引 (num_elements, 3) ----
    @property
    def element_indices(self) -> np.ndarray:
        return self._mesh.t.T

    # ---- 顶点坐标 (num_vertices, 2) ----
    @property
    def vertex_positions(self) -> np.ndarray:
        return self._mesh.p.T

    # ---- 网格边（相邻顶点对索引）----
    @property
    def mesh_edges(self) -> np.ndarray:
        return self._mesh.facets

    # ---- 每个单元的相邻单元（排除边界）----
    @property
    def element_neighbors(self) -> np.ndarray:
        return self._mesh.f2t[:, self._mesh.f2t[1] != -1]

    # ---- 可视化用的计算域边界框 ----
    @property
    def plot_boundary(self):
        return self._plot_boundary

    # ---- 额外的 plotly 渲染图层 ----
    def additional_plots(self) -> Dict:
        return self.fem_problem.additional_plots_from_mesh(self._mesh)


class FEMProblemCircularQueue:
    """Circular buffer of Helmholtz instances for training generalization."""

    def __init__(
        self,
        *,
        fem_config: Dict[Union[str, int], Any],
        random_state: np.random.RandomState = np.random.RandomState(),
    ):
        self._fem_config = fem_config
        self._random_state = random_state

        num_pdes = fem_config.get("num_pdes", 100)
        self._use_buffer = num_pdes is not None and num_pdes > 0
        num_pdes = num_pdes if self._use_buffer else 1

        self._index_sampler = IndexSampler(num_pdes, random_state=self._random_state)
        self._fem_problems: List[Optional[FEMProblemWrapper]] = [None for _ in range(num_pdes)]

        pde_config = fem_config.get(fem_config.get("pde_type", "helmholtz"), {})
        self._pde_features = {
            "element_features": [
                name for name, include in pde_config.get("element_features", {}).items() if include
            ],
        }

    def next(self) -> FEMProblemWrapper:
        return self._next_from_idx(pde_idx=self._index_sampler.next())

    def _next_from_idx(self, pde_idx: int) -> FEMProblemWrapper:
        if (not self._use_buffer) or self._fem_problems[pde_idx] is None:
            new_seed = self._random_state.randint(0, 2**31)
            new_problem = create_helmholtz_problem(
                fem_config=self._fem_config,
                random_state=np.random.RandomState(seed=new_seed),
            )
            self._fem_problems[pde_idx] = FEMProblemWrapper(
                fem_config=self._fem_config,
                fem_problem=new_problem,
                pde_features=self._pde_features,
            )
        self._fem_problems[pde_idx].reset()
        return self._fem_problems[pde_idx]

    # PDE 提供的单元特征个数
    @property
    def num_pde_element_features(self) -> int:
        return len(self._pde_features["element_features"])
