ScalarField

@beartype
class ScalarField(AbstractField):
    """A scalar field defined on a mesh.

    Each degree of freedom holds a single scalar value.

    """

    def __init__(self, *args, **kwargs):
        """ScalarField init to check element type."""
        super().__init__(*args, **kwargs)
        if isinstance(self._element_type, skfem.ElementVector):
            raise TypeError("The element type of a scalar field cannot be vector.")

    @property
    def dof_locations(self: Self) -> NDArray:
        """Return the coordinates of all degrees of freedom.

        Returns:
            Array of DOF coordinates with shape ``(n_dofs, 3)``.

        """
        return self.skfem_basis.doflocs.T

    @property
    def element_type(self: Self) -> skfem.Element:
        """Return the finite element type.

        Returns:
            A scikit-fem element instance. Guaranteed not to be an ``ElementVector``.

        """
        return self._element_type

    @property
    def value(self: Self) -> NDArray:
        """Return the field values at the degrees of freedom.

        Returns:
            A 1D array of scalar values with shape ``(n_dofs,)``.

        """
        return self._value

    @value.setter
    def value(
        self: Self,
        value: Number | list | tuple | NDArray | Callable,
    ):
        """Set the field values.

        Args:
            value: New field values. Accepts a scalar (constant field),
                an arrayan like of shape ``(n_dofs,)``,
                or a callable that maps degree of freedom coordinates to a 1D array.

        Raises:
            ValueError:
                - If the array like shape does not match the number of DOFs.
                - If the callable output array shape does not match the number of DOFs.
            TypeError: If the output of callable is not a NDArray.

        """
        if isinstance(value, (int, float)):
            self._value = np.tile(value, self.dof_locations.shape[0])
        elif isinstance(value, (list, tuple)):
            if len(value) != self.dof_locations.shape[0]:
                raise ValueError("Value of a scalar field cannot be a vector.")
            else:
                self._value = np.array(value)
        elif isinstance(value, np.ndarray):
            if value.shape != (self.dof_locations.shape[0],):
                raise ValueError(
                    "The value array of the scalar field must be a 1D array "
                    f"of length {self.dof_locations.shape[0]}. "
                    f"Instead, an array of shape {value.shape} was passed."
                )
            else:
                self._value = value
        elif isinstance(value, Callable):
            out_arr = value(self.dof_locations)
            if not isinstance(out_arr, np.ndarray):
                raise TypeError(
                    "The output of the function used to set the value of the field "
                    "must be a numpy array. "
                    f"Instead, {type(out_arr).__name__} was obtained."
                )
            elif out_arr.shape != (self.dof_locations.shape[0],):
                raise ValueError(
                    "The output of the function used to set the value of the "
                    "scalar field should be a 1D array of length "
                    f"{self.dof_locations.shape[0]}. "
                    f"Instead, {out_arr.shape} shaped array was obtained."
                )
            self._value = out_arr

    def set_subregion_value(
        self: Self, name: str, value: Number | list | tuple | NDArray | Callable
    ):
        """Set field values on a named subregion.

        Args:
            name: Name of the subregion to modify.
            value: New values for the subregion DOFs. Accepts the same types
                as the ``value`` setter.

        Raises:
            RuntimeError: If no subregions are defined on the mesh.
            ValueError: Same as ``value`` setter or the subregion name does not exist.
            TypeError: Same as ``value`` setter.

        """
        if not self.mesh.subregions:
            raise RuntimeError("No subregion defined!")
        else:
            subregion_names = self.mesh.subregions.keys()
        if name not in subregion_names:
            raise ValueError(
                f"The subregion {name} does not exist. "
                f"Please select one from {subregion_names}."
            )

        dof_ids = self.get_subregion_dof_index(name)
        subregion_dofs = self.dof_locations[dof_ids]

        if isinstance(value, (int, float)):
            self._value[dof_ids] = np.tile(value, subregion_dofs.shape[0])
        elif isinstance(value, (list, tuple)):
            if len(value) != subregion_dofs.shape[0]:
                raise ValueError("Value of a scalar field cannot be a vector.")
            else:
                self._value[dof_ids] = np.array(value)
        elif isinstance(value, np.ndarray):
            if value.shape != (subregion_dofs.shape[0],):
                raise ValueError(
                    "The value array of the scalar field must be a 1D array "
                    f"of length {subregion_dofs.shape[0]}. "
                    f"Instead, an array of shape {value.shape} was passed."
                )
            else:
                self._value[dof_ids] = value
        elif isinstance(value, Callable):
            out_arr = value(subregion_dofs)
            if not isinstance(out_arr, np.ndarray):
                raise TypeError(
                    "The output of the function used to set the value of the field "
                    "must be a numpy array. "
                    f"Instead, {type(out_arr).__name__} was obtained."
                )
            elif out_arr.shape != (subregion_dofs.shape[0],):
                raise ValueError(
                    "The output of the function used to set the value of the "
                    "scalar field should be a 1D array of length "
                    f"{subregion_dofs.shape[0]}. "
                    f"Instead, {out_arr.shape} shaped array was obtained."
                )
            self._value[dof_ids] = out_arr

    def to_skfem(self: Self) -> skfem.DiscreteField:
        """Convert the field to a scikit-fem discrete field.

        Returns:
            A scikit-fem ``DiscreteField`` interpolated from the field values.

        """
        return self.skfem_basis.interpolate(self.value)

    def to_pyvista(self: Self) -> pv.UnstructuredGrid:
        """Convert the field to a PyVista UnstructuredGrid object.

        Returns:
            PyVista ``UnstructuredGrid`` with ``value`` data.

        Notes:
            Higher-order elements are projected to linear P1 elements for
            visualisation.

        """
        if isinstance(
            self.element_type,
            (
                # Check for DOFs on the nodes or cells.
                skfem.ElementLineP1,
                skfem.ElementTriP1,
                skfem.ElementTetP1,
                skfem.ElementQuad1,
                skfem.ElementHex1,
            ),
        ):
            mesh_pv_main: pv.UnstructuredGrid = self.mesh.to_pyvista()

            if mesh_pv_main.points.shape == self.dof_locations.shape and np.allclose(
                mesh_pv_main.points, self.dof_locations
            ):
                final_value = self.value
            else:
                tree = KDTree(mesh_pv_main.points)
                dist, new_point_indices = tree.query(self.dof_locations)
                if not np.allclose(np.zeros_like(dist), dist):
                    raise RuntimeError(
                        "Unable to create new indices for scikit-fem mesh points."
                    )
                final_value = np.zeros(mesh_pv_main.points.shape[0])
                final_value[new_point_indices] = self.value

            mesh_pv_main.point_data["value"] = final_value
            mesh_pv_main.set_active_scalars("value", "point")
            return mesh_pv_main
        elif isinstance(
            self.element_type,
            (
                # Check for DOFs on the nodes or cells.
                skfem.ElementLineP0,
                skfem.ElementTriP0,
                skfem.ElementTetP0,
                skfem.ElementQuad0,
                skfem.ElementHex0,
            ),
        ):
            mesh_pv_main: pv.UnstructuredGrid = self.mesh.to_pyvista()

            if (
                mesh_pv_main.cell_centers().points.shape == self.dof_locations.shape
                and np.allclose(mesh_pv_main.cell_centers().points, self.dof_locations)
            ):
                final_value = self.value
            else:
                tree = KDTree(mesh_pv_main.cell_centers().points)
                dist, new_point_indices = tree.query(self.dof_locations)
                if not np.allclose(np.zeros_like(dist), dist):
                    raise RuntimeError(
                        "Unable to create new indices for dof_locations."
                    )
                final_value = np.zeros(mesh_pv_main.cell_centers().points.shape[0])
                final_value[new_point_indices] = self.value
            mesh_pv_main.cell_data["value"] = final_value
            mesh_pv_main.set_active_scalars("value", "cell")
            return mesh_pv_main
        else:
            skfem_mesh = self.mesh.to_skfem()

            if isinstance(skfem_mesh, skfem.mesh.MeshLine1):
                new_element = skfem.ElementLineP1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshTri1):
                new_element = skfem.ElementTriP1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshTet1):
                new_element = skfem.ElementTetP1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshQuad1):
                new_element = skfem.ElementQuad1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshHex1):
                new_element = skfem.ElementHex1()
            else:
                raise TypeError(f"Unrecognised mesh type {skfem_mesh}.")

            new_value = self.skfem_basis.with_element(new_element).project(
                self.to_skfem()
            )
            new_field = ScalarField(self.mesh, new_element, new_value)
            return new_field.to_pyvista()

    def plot(
        self: Self, subregions: list[str] | None = None, cmap: str = "viridis", **kwargs
    ):
        """Plot the scalar field using PyVista.

        Args:
            subregions: Optional list defining subregions to plot.
            cmap: Matplotlib colormap name. Defaults to ``"viridis"``.
            **kwargs: Additional keyword arguments forwarded to
                ``pyvista.Plotter.show``.

        """
        plot_pv_obj: pv.UnstructuredGrid = self.to_pyvista()
        pl = pv.Plotter()
        if (
            subregions
            and self.mesh.subregions
            and not set(subregions).issubset(set(self.mesh.subregions))
        ):
            raise ValueError(
                f"Subregions to be plotted {subregions} "
                f"is not a subset of mesh subregions {self.mesh.subregions.keys()}."
            )
        elif subregions and self.mesh.subregions is None:
            raise ValueError(
                f"Cannot plot subregions {subregions} "
                "since no subregions are defined on the mesh."
            )
        elif subregions:
            for subregion in subregions:
                pl.add_mesh(
                    plot_pv_obj.threshold(value=True, scalars=subregion), cmap=cmap
                )
        else:
            pl.add_mesh_clip_plane(
                plot_pv_obj, outline_opacity=0.3, cmap=cmap, invert=True
            )
        pl.show_axes_all()
        pl.show(**kwargs)

    def __str__(self: Self) -> str:
        """Return a concise string representation.

        Returns:
            Description the field type, element type, subregions, and value range.

        """
        element_type = type(self.element_type).__name__
        s = f"{self.__class__.__name__}(elements={element_type}"
        if self.mesh.subregions:
            s += f", subregions={[s for s in self.mesh.subregions]}"
        s += f", range=[{self.value.min()}, {self.value.max()}])"

        return s

    def __repr__(self: Self) -> str:
        """Return an unambiguous string representation.

        Returns:
            Same as ``__str__``.

        """
        return self.__str__()

    def __eq__(self: Self, other: object) -> bool:
        """Compare two scalar fields for equality.

        Returns:
            ``True`` for the same mesh, element type, and values (numpy tolerance).

        """
        if not isinstance(other, self.__class__):
            return False
        return (
            self.mesh == other.mesh
            and type(self.element_type).__name__ == type(other.element_type).__name__
            and np.allclose(self.value, other.value)
        )

    def __call__(self: Self, point: list | tuple | NDArray) -> NDArray:
        """Evaluate the scalar field at one or more points.

        Args:
            point: A single 3D point (length-3 list, tuple, or 1D array) or a
                2D array of points with shape ``(n_points, 3)``.

        Returns:
            Interpolated scalar values at the requested points.

        Raises:
            ValueError: If the point coordinates are not 3D or have an
                unsupported shape. A value error is also raised from
                scikit-fem if the point is outside the mesh bounds.

        """
        if isinstance(point, list | tuple):
            if len(point) != 3:
                raise ValueError(
                    "The value of point for which to calculate the value of field "
                    f"needs to be a length 3 {type(point).__name__}. instead {len(point)=}."
                )
            else:
                point = np.array([point]).T
        elif isinstance(point, np.ndarray):
            if point.ndim == 1 and len(point) != 3:
                raise ValueError(
                    "The length of 1D point value array must be 3. "
                    f"instead {len(point)=}."
                )
            elif point.ndim == 2 and point.shape[-1] != 3:
                raise ValueError(
                    "The last dimension of 2D point value array must be 3. "
                    f"instead it is {point.shape[-1]}."
                )
            elif point.ndim > 2:
                raise ValueError(
                    "The point value array must be 1D of length 3 "
                    "or 2D with last dimension as 3. "
                    f"Instead {point.shape=}."
                )
            else:
                point = np.array([point]).T if point.ndim == 1 else point.T

        calc_fun = self.skfem_basis.interpolator(self.value)
        return calc_fun(point)

    def get_subregion_dof_index(self: Self, name: str) -> NDArray:
        """Return DOF indices belonging to a named subregion or the field boundary.

        Args:
            name: Name of the subregion or boundary.

        Returns:
            Array of DOF indices.

        Raises:
            ValueError: If the subregion is not found in either subdomains or
                boundaries.

        """
        skfem_elem_dofs = self.skfem_basis.element_dofs.T
        skfem_mesh = self.mesh.to_skfem()

        if name.lower() == "boundary":
            return self.skfem_basis.get_dofs().flatten()
        elif skfem_mesh.subdomains is not None and name in skfem_mesh.subdomains:
            dof_indices = np.unique(skfem_elem_dofs[skfem_mesh.subdomains[name]])
            return dof_indices
        elif skfem_mesh.boundaries is not None and name in skfem_mesh.boundaries:
            dof_indices = self.skfem_basis.get_dofs(name).flatten()
            return dof_indices
        else:
            raise ValueError(f"Unable to find subregion {name}.")