"""Classes representing (perturbed) droplets in various dimensions.

The classes differ in what features of a droplet they track. In the simplest case, only
the position and radius of a spherical droplet are stored. Other classes additionally
keep track of the interfacial width or shape perturbations (of various degrees).

.. autosummary::
   :nosignatures:

   SphericalDroplet
   DiffuseDroplet
   PerturbedDroplet2D
   PerturbedDroplet3D
   PerturbedDroplet3DAxisSym


Inheritance structure of the classes:

.. inheritance-diagram:: droplets.droplets
   :parts: 1

The details of the classes are explained below:

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""

from __future__ import annotations

import logging
import math
import warnings
from abc import ABCMeta, abstractmethod
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypeVar

import numpy as np
from numba.extending import register_jitable
from numpy.lib.recfunctions import structured_to_unstructured
from scipy import integrate
from typing_extensions import Self

from pde.fields import ScalarField
from pde.tools.cuboid import Cuboid
from pde.tools.plotting import PlotReference, plot_on_axes

from .tools import spherical
from .tools.misc import enable_scalar_args

if TYPE_CHECKING:
    from collections.abc import Callable

    from numpy.typing import DTypeLike

    from pde.grids.base import GridBase

    from .tools.typing import RealArray

_logger = logging.getLogger(__name__)
""":class:`logging.Logger`: Logger instance."""


TDroplet = TypeVar("TDroplet", bound="DropletBase")
DTypeList = list[tuple[str, type[Any]] | tuple[str, type[Any], tuple[int, ...]]]


def get_dtype_field_size(dtype: DTypeLike, field_name: str) -> int:
    """Return the number of elements in a field of structured numpy array.

    Args:
        dtype (list):
            The dtype of the numpy array
        field_name (str):
            The name of the field that needs to be checkedrm
    """
    shape = dtype.fields[field_name][0].shape  # type: ignore
    return np.prod(shape) if shape else 1  # type: ignore


def iterate_in_pairs(it, fill=0):
    """Return consecutive pairs from an iterator.

    For instance, `list(pair_iterator('ABCDE', fill='Z'))` returns
    `[('A', 'B'), ('C', 'D'), ('E', 'Z')]`

    Args:
        it (iterator):
            The iterator
        fill:
            The value returned for the second part of the last returned tuple if the
            length of the iterator is odd

    Returns:
        This is a generator function that yields pairs of items of the iterator
    """
    it = iter(it)
    while True:
        # obtain first item of tuple
        try:
            first = next(it)
        except StopIteration:
            break
        # obtain second item of tuple
        try:
            yield first, next(it)
        except StopIteration:
            yield first, fill
            break


class DropletBase:
    """Represents a generic droplet.

    The data associated with a droplet is stored in structured array. Consequently, the
    `dtype` of the array determines what information the droplet class stores.
    """

    _subclasses: dict[str, DropletBase] = {}  # collect all inheriting classes

    __slots__ = ["data"]

    data: np.recarray  # all information about the droplet in a record array

    _merge_data: Callable[[RealArray, RealArray, RealArray], None]
    """Private method for merging droplet data, created by __init_subclass__"""

    @classmethod
    def from_data(cls, data: np.recarray) -> Self:
        """Create droplet class from a given data.

        Args:
            data (:class:`numpy.recarray`):
                The data array used to initialize the droplet
        """
        # note that we do not call the __init__ method of the class, since we do
        # not need to invent the dtype and set all the attributes. We here
        # simply assume that the given data is sane
        obj = cls.__new__(cls)
        obj.data = data
        return obj

    @classmethod
    def from_droplet(cls, droplet: DropletBase, **kwargs) -> Self:
        r"""Return a droplet with data taken from `droplet`

        Args:
            droplet (:class:`DropletBase`):
                Another droplet from which data is copied. This does not be the same
                type of droplet
            \**kwargs:
                Additional arguments that can overwrite data in `droplet` or set
                additional arguments for the current class
        """
        args = droplet._args
        args.update(kwargs)
        return cls(**args)

    @classmethod
    @abstractmethod
    def get_dtype(cls, **kwargs) -> DTypeList:
        """Determine the dtype representing this droplet class.

        Returns:
            :class:`numpy.dtype`:
                The (structured) dtype associated with this class
        """
        ...

    def _init_data(self, **kwargs) -> None:
        """Initializes the `data` attribute if it is not present.

        Args:
            **kwargs:
                Arguments used to determine the dtype of the data array
        """
        if not hasattr(self, "data"):
            dtype = self.get_dtype(**kwargs)

            # We need to create an empty record with the correct data type. Note
            # that the conversion np.recarray(1)[0] turns the array into a
            # scalar type (instance of numpy.record) that contains the
            # structured data. This conversion is necessary for numba to operate
            # on the data.
            self.data = np.recarray(1, dtype=dtype)[0]

    def __init_subclass__(cls, **kwargs):
        """Modify subclasses of this base class."""
        super().__init_subclass__(**kwargs)

        # register all subclasses to reconstruct them later
        if cls is not DropletBase:
            if cls.__name__ in cls._subclasses:
                warnings.warn(f"Redefining class {cls.__name__}", stacklevel=2)
            cls._subclasses[cls.__name__] = cls

        # create a staticmethod for merging droplet data
        cls._merge_data = staticmethod(cls._make_merge_data())

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return NotImplemented
        return np.allclose(
            self._data_array, other._data_array, rtol=0, atol=0, equal_nan=True
        )

    def check_data(self):
        """Method that checks the validity and consistency of self.data."""

    @property
    def _args(self):
        return {key: self.data[key] for key in self.data.dtype.names}

    def __str__(self):
        arg_list = [f"{key}={value}" for key, value in self._args.items()]
        return f"{self.__class__.__name__}({', '.join(arg_list)})"

    __repr__ = __str__

    @property
    def _data_array(self) -> RealArray:
        """:class:`~numpy.ndarray`: the data of the droplet in an unstructured array."""
        return structured_to_unstructured(self.data)

    def copy(self, **kwargs) -> Self:
        r"""Return a copy of the current droplet.

        Args:
            \**kwargs:
                Additional arguments an be used to set data of the returned droplet.
        """
        if kwargs:
            return self.from_droplet(self, **kwargs)
        return self.from_data(self.data.copy())

    @classmethod
    def _make_merge_data(cls) -> Callable[[RealArray, RealArray, RealArray], None]:
        """Factory for a function that merges the data of two droplets."""
        raise NotImplementedError

    def merge(self, other: DropletBase, *, inplace: bool = False) -> Self:
        """Merge two droplets into one."""
        if inplace:
            self._merge_data(self.data, other.data, out=self.data)  # type: ignore
            return self
        result = np.record(np.zeros_like(self.data))
        self._merge_data(self.data, other.data, out=result)  # type: ignore
        return self.__class__.from_data(result)  # type: ignore

    @property
    def data_bounds(self) -> tuple[RealArray, RealArray]:
        """tuple: lower and upper bounds on the parameters"""
        num = len(self._data_array)
        return np.full(num, -np.inf), np.full(num, np.inf)


class SphericalDroplet(DropletBase):
    """Represents a single, spherical droplet."""

    __slots__ = ["data"]

    def __init__(self, position: RealArray, radius: float):
        r"""
        Args:
            position (:class:`~numpy.ndarray`):
                Position of the droplet center
            radius (float):
                Radius of the droplet
        """
        self._init_data(position=position)

        self.position = position
        self.radius = radius
        self.check_data()

    def check_data(self):
        """Method that checks the validity and consistency of self.data."""
        if self.radius < 0:
            msg = "Radius must be positive"
            raise ValueError(msg)

    @classmethod
    def get_dtype(cls, **kwargs) -> DTypeList:
        """Determine the dtype representing this droplet class.

        Args:
            position (:class:`~numpy.ndarray`):
                The position vector of the droplet, determining space dimension.

        Returns:
            :class:`numpy.dtype`: the (structured) dtype associated with this class
        """
        position = np.atleast_1d(kwargs.pop("position"))
        if kwargs:
            msg = f"Leftover keyword arguments: {kwargs}"
            raise ValueError(msg)
        dim = len(position)
        return [("position", float, (dim,)), ("radius", float)]

    @property
    def dim(self) -> int:
        """int: the spatial dimension this droplet is embedded in"""
        return get_dtype_field_size(self.data.dtype, "position")

    @property
    def data_bounds(self) -> tuple[RealArray, RealArray]:
        """tuple: lower and upper bounds on the parameters"""
        l, h = super().data_bounds
        l[self.dim] = 0  # radius must be non-negative
        return l, h

    @classmethod
    def from_volume(cls, position: RealArray, volume: float, **kwargs):
        r"""Construct a droplet from given volume instead of radius.

        Args:
            position (:class:`~numpy.ndarray`):
                Center of the droplet
            volume (float):
                Volume of the droplet
            \**kwargs:
                Additional arguments are forwarded to the class initializer. This can
                for instance be used to set the interfacial width.
        """
        dim = len(np.array(position, np.double, ndmin=1))
        radius = spherical.radius_from_volume(volume, dim)
        return cls(position, radius, **kwargs)

    @property
    def position(self) -> RealArray:
        """:class:`~numpy.ndarray`: the position of the droplet."""
        return self.data["position"]

    @position.setter
    def position(self, value: RealArray) -> None:
        value = np.asanyarray(value)
        if len(value) != self.dim:
            msg = f"The dimension of the position must be {self.dim}"
            raise ValueError(msg)
        self.data["position"] = value

    @property
    def radius(self) -> float:
        """float: the radius of the droplet"""
        return float(self.data["radius"])

    @radius.setter
    def radius(self, value: float) -> None:
        self.data["radius"] = value
        self.check_data()

    @property
    def volume(self) -> float:
        """float: volume of the droplet"""
        return spherical.volume_from_radius(self.radius, self.dim)

    @volume.setter
    def volume(self, volume: float) -> None:
        """Set the radius from a supplied volume."""
        self.radius = spherical.radius_from_volume(volume, self.dim)

    @property
    def surface_area(self) -> float:
        """float: surface area of the droplet"""
        return spherical.surface_from_radius(self.radius, self.dim)

    @property
    def bbox(self) -> Cuboid:
        """:class:`~pde.tools.cuboid.Cuboid`: bounding box of the droplet."""
        return Cuboid.from_points(
            self.position - self.radius, self.position + self.radius
        )

    def overlaps(self, other: SphericalDroplet, grid: GridBase | None = None) -> bool:
        """Determine whether another droplet overlaps with this one.

        Note that this function so far only compares the distances of the droplets to
        their radii, which does not respect perturbed droplets correctly.

        Args:
            other (:class:`~droplets.droplets.SphericalDroplet`):
                instance of the other droplet
            grid (:class:`~pde.grids.base.GridBase`):
                grid that determines how distances are measured, which is for instance
                important to respect periodic boundary conditions. If omitted, an
                Euclidean distance is assumed.

        Returns:
            bool: whether the droplets overlap or not
        """
        if grid is None:
            distance = float(np.linalg.norm(self.position - other.position))
        else:
            distance = grid.distance(self.position, other.position, coords="cartesian")  # type: ignore
        return distance < self.radius + other.radius

    @classmethod
    def _make_merge_data(cls) -> Callable[[RealArray, RealArray, RealArray], None]:
        """Factory for a function that merges the data of two droplets."""
        radius_from_volume = spherical.make_radius_from_volume_nd_compiled()
        volume_from_radius = spherical.make_volume_from_radius_nd_compiled()

        @register_jitable
        def merge_data(drop1: RealArray, drop2: RealArray, out: RealArray) -> None:
            """Merge the data of two droplets."""
            dim = len(drop1.position)  # type: ignore

            V1 = volume_from_radius(drop1.radius, dim)  # type: ignore
            V2 = volume_from_radius(drop2.radius, dim)  # type: ignore
            volume = V1 + V2
            out.radius = radius_from_volume(volume, dim)  # type: ignore

            # adjust droplet position
            out.position[...] = (V1 * drop1.position + V2 * drop2.position) / volume  # type: ignore

        return merge_data  # type: ignore

    @enable_scalar_args
    def interface_position(self, *args) -> RealArray:
        r"""Calculates the position of the interface of the droplet.

        Args:
            *args (float or :class:`~numpy.ndarray`):
                The angles identifying the interface points. For 2d droplets, this is
                simply the angle in polar coordinates. For 3d droplets, both the
                azimuthal angle θ (in :math:`[0, \pi]`) and the polar angle φ (in
                :math:`[0, 2\pi]`) need to be specified.

        Returns:
            :class:`~numpy.ndarray`: An array with the coordinates of the interfacial
            points associated with each angle given by `φ`.

        Raises:
            ValueError: If the dimension of the space is not 2
        """
        if self.dim != len(args) + 1:
            msg = f"Interfacial position requires {self.dim - 1} angles"
            raise ValueError(msg)

        if self.dim == 2:
            # spherical droplet in two dimensions
            φ = args[0]
            pos = self.radius * np.transpose([np.cos(φ), np.sin(φ)])

        elif self.dim == 3:
            # spherical droplet in three dimensions
            θ, φ = args[0], args[1]
            r = np.full_like(θ, self.radius)
            pos = spherical.points_spherical_to_cartesian(np.c_[r, θ, φ])  # type: ignore

        else:
            msg = f"Cannot calculate {self.dim}d position"
            raise NotImplementedError(msg)

        # shift the droplet center
        return self.position[None, :] + pos

    @property
    def interface_curvature(self) -> float:
        """float: the mean curvature of the interface of the droplet"""
        return 1 / self.radius

    def _get_phase_field(self, grid: GridBase, dtype: DTypeLike = float) -> RealArray:
        """Creates an image of the droplet on the `grid`

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                The grid used for discretizing the droplet phase field
            dtype (class:`numpy.dtype`):
                The numpy data type defining the type of the returned image.
                If `dtype == np.bool`, a binary representation is returned.

        Returns:
            :class:`~numpy.ndarray`: An array with data values representing the droplet
            phase field at support points of the `grid`.
        """
        if self.dim != grid.dim:
            msg = (
                f"Droplet (dimension {self.dim}) incompatible with grid (dimension "
                f"{grid.dim})"
            )
            raise ValueError(msg)

        # calculate distances from droplet center
        dist = spherical.polar_coordinates(grid, origin=self.position, ret_angle=False)
        return (dist < self.radius).astype(dtype)

    def get_phase_field(
        self,
        grid: GridBase,
        *,
        vmin: float = 0,
        vmax: float = 1,
        label: str | None = None,
    ) -> ScalarField:
        """Creates an image of the droplet on the `grid`

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                The grid used for discretizing the droplet phase field
            vmin (float):
                Minimal value the phase field will attain (far away from droplet)
            vmax (float):
                Maximal value the phase field will attain (inside the droplet)
            label (str):
                The label associated with the returned scalar field

        Returns:
            :class:`~pde.fields.ScalarField`: A scalar field
            representing the droplet
        """
        data = self._get_phase_field(grid)
        data = vmin + (vmax - vmin) * data  # scale data
        return ScalarField(grid, data=data, label=label)

    def get_triangulation(self, resolution: float = 1) -> dict[str, Any]:
        """Obtain a triangulated shape of the droplet surface.

        Args:
            resolution (float):
                The length of a typical triangulation element. This affects the
                resolution of the triangulation.

        Returns:
            dict: A dictionary containing information about the triangulation. The exact
            details depend on the dimension of the problem.
        """
        if self.dim == 2:
            num = max(3, int(np.ceil(self.surface_area / resolution)))
            angles = np.linspace(0, 2 * np.pi, num + 1, endpoint=True)
            vertices = self.interface_position(angles)
            lines = np.c_[np.arange(num), np.arange(1, num + 1) % num]
            return {"vertices": vertices, "lines": lines}

        if self.dim == 3:
            # estimate the number of triangles covering the surface
            try:
                surface_area = self.surface_area
            except (NotImplementedError, AttributeError):
                # estimate surface area from 3d spherical droplet
                surface_area = spherical.surface_from_radius(self.radius, dim=3)
            num_est = (4 * surface_area) / (np.sqrt(3) * resolution**2)
            tri = triangulated_spheres.get_triangulation(num_est)

            φ, θ = tri["angles"][:, 0], tri["angles"][:, 1]
            return {
                "vertices": self.interface_position(θ, φ),
                "triangles": tri["cells"],
            }

        msg = f"Triangulation not implemented for {self.dim}d"
        raise NotImplementedError(msg)

    def _get_mpl_patch(self, dim=None, **kwargs):
        """Return the patch representing the droplet for plotting.

        Args:
            dim (int, optional):
                The dimension in which the data is plotted. If omitted, the actual
                physical dimension is assumed.
            **kwargs:
                Additional keyword arguments are passed to
                :class:`matplotlib.patches.Circle`, which creates the patch that
                represents the droplet. For instance, to only draw the outlines of the
                droplets, you may need to supply `fill=False`.

        Returns:
            :class:`~matplotlib.patches.Circle`: The patch representing the droplet
        """
        import matplotlib as mpl

        if dim is None:
            dim = self.dim

        if dim != 2:
            msg = "Plotting is only implemented in 2d"
            raise NotImplementedError(msg)

        if self.dim == 1:
            position = (self.position[0], 0)
        else:
            position = self.position[:dim]

        return mpl.patches.Circle(position, self.radius, **kwargs)

    @plot_on_axes()
    def plot(self, ax, value: Callable | None = None, **kwargs) -> PlotReference:
        """Plot the droplet.

        Args:
            {PLOT_ARGS}
            value (callable):
                Sets the color of the droplet. This could be either `a matplotlib color
                <https://matplotlib.org/stable/tutorials/colors/colors.html>`_ or a
                function that takes the droplet instance and returns a color in which
                this droplet is drawn. If given, it overwrites the `color` argument.
            **kwargs:
                Additional keyword arguments are passed to the class that creates the
                patch that represents the droplet. For instance, to only draw the
                outlines of the droplets, you may need to supply `fill=False`.

        Returns:
            :class:`~pde.tools.plotting.PlotReference`: Information about the plot
        """
        # create the artist representing the droplet
        artist = self._get_mpl_patch(**kwargs)
        # add artist to axis and return all information in a plot reference
        ax.add_artist(artist)
        return PlotReference(ax, artist)


class DiffuseDroplet(SphericalDroplet):
    """Represents a single, spherical droplet with a diffuse interface."""

    __slots__ = ["data"]

    def __init__(
        self,
        position: RealArray,
        radius: float,
        interface_width: float | None = None,
    ):
        """
        Args:
            position (:class:`~numpy.ndarray`):
                Position of the droplet center
            radius (float):
                Radius of the droplet
            interface_width (float, optional):
                Width of the interface
        """
        self._init_data(position=position)
        super().__init__(position=position, radius=radius)
        self.interface_width = interface_width

    @property
    def data_bounds(self) -> tuple[RealArray, RealArray]:
        """tuple: lower and upper bounds on the parameters"""
        l, h = super().data_bounds
        l[self.dim + 1] = 0  # interface width must be non-negative
        return l, h

    @classmethod
    def get_dtype(cls, **kwargs) -> DTypeList:
        """Determine the dtype representing this droplet class.

        Args:
            position (:class:`~numpy.ndarray`):
                The position vector of the droplet, determining the space dimension.

        Returns:
            :class:`numpy.dtype`: the (structured) dtype associated with this class
        """
        dtype = super().get_dtype(**kwargs)
        return [*dtype, ("interface_width", float)]

    @classmethod
    def _make_merge_data(cls) -> Callable[[RealArray, RealArray, RealArray], None]:
        """Factory for a function that merges the data of two droplets."""
        parent_merge = super()._make_merge_data()

        @register_jitable
        def merge_data(drop1: RealArray, drop2: RealArray, out: RealArray) -> None:
            """Merge the data of two droplets."""
            parent_merge(drop1, drop2, out)
            out.interface_width = (drop1.interface_width + drop2.interface_width) / 2  # type: ignore

        return merge_data  # type: ignore

    @property
    def interface_width(self) -> float | None:
        """float: the width of the interface of this droplet"""
        if np.isnan(self.data["interface_width"]):
            return None
        return float(self.data["interface_width"])

    @interface_width.setter
    def interface_width(self, value: float | None) -> None:
        if value is None:
            self.data["interface_width"] = math.nan
        elif value < 0:
            msg = "Interface width must not be negative"
            raise ValueError(msg)
        else:
            self.data["interface_width"] = value
        self.check_data()

    def _get_phase_field(self, grid: GridBase, dtype: DTypeLike = float) -> RealArray:
        """Creates an image of the droplet on the `grid`

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                The grid used for discretizing the droplet phase field
            dtype (:class:`numpy.dtype`):
                The numpy data type defining the type of the returned image.
                If `dtype == np.bool`, a binary representation is returned.

        Returns:
            :class:`~numpy.ndarray`: An array with data values representing the droplet
            phase field at support points of the `grid`.
        """
        if self.dim != grid.dim:
            msg = (
                f"Droplet (dimension {self.dim}) incompatible with grid (dimension "
                f"{grid.dim})"
            )
            raise ValueError(msg)

        if self.interface_width is None:
            interface_width = grid.typical_discretization
        else:
            interface_width = self.interface_width

        # calculate distances from droplet center
        dist = spherical.polar_coordinates(grid, origin=self.position, ret_angle=False)

        # make the image
        if interface_width == 0 or np.issubdtype(dtype, bool):
            result = dist < self.radius
        else:
            result = 0.5 + 0.5 * np.tanh(2 * (self.radius - dist) / interface_width)

        return result.astype(dtype)


class PerturbedDropletBase(DiffuseDroplet, metaclass=ABCMeta):
    """Represents a single droplet with a perturbed shape.

    This acts as an abstract class for which member functions need to specified
    depending on dimensionality.
    """

    __slots__ = ["data"]

    def __init__(
        self,
        position: RealArray,
        radius: float,
        interface_width: float | None = None,
        amplitudes: RealArray | None = None,
    ):
        """
        Args:
            position (:class:`~numpy.ndarray`):
                Position of the droplet center
            radius (float):
                Radius of the droplet
            interface_width (float, optional):
                Width of the interface
            amplitudes (:class:`~numpy.ndarray`):
                The amplitudes of the perturbations
        """
        self._init_data(position=position, amplitudes=amplitudes)
        super().__init__(
            position=position, radius=radius, interface_width=interface_width
        )

        self.amplitudes = amplitudes

        if len(self.position) != self.__class__.dim:
            msg = f"Space dimension must be {self.__class__.dim}"
            raise ValueError(msg)

    @classmethod
    def get_dtype(cls, **kwargs) -> DTypeList:
        """Determine the dtype representing this droplet class.

        Args:
            position (:class:`~numpy.ndarray`):
                The position vector of the droplet. This is used to determine the space
                dimension.
            amplitudes (:class:`~numpy.ndarray`):
                The perturbation amplitudes used to determine their size

        Returns:
            :class:`numpy.dtype`: the (structured) dtype associated with this class
        """
        # extract data
        amplitudes = kwargs.pop("amplitudes")
        if amplitudes is None:
            modes = 0
        else:
            modes = len(amplitudes)

        # create dtype
        dtype = super().get_dtype(**kwargs)
        return [*dtype, ("amplitudes", float, (modes,))]

    @property
    def data_bounds(self) -> tuple[RealArray, RealArray]:
        """tuple: lower and upper bounds on the parameters"""
        l, h = super().data_bounds
        n = self.dim + 2
        # relative perturbation amplitudes must be between [-1, 1]
        l[n : n + self.modes] = -1
        h[n : n + self.modes] = 1
        return l, h

    @property
    def modes(self) -> int:
        """int: number of perturbation modes"""
        shape = self.data.dtype.fields["amplitudes"][0].shape  # type: ignore
        return int(shape[0]) if shape else 1

    @property
    def amplitudes(self) -> RealArray:
        """:class:`~numpy.ndarray`: the perturbation amplitudes."""
        return np.atleast_1d(self.data["amplitudes"])

    @amplitudes.setter
    def amplitudes(self, value: RealArray | None = None) -> None:
        if value is None:
            if self.modes != 0:
                msg = "Require values for amplitudes"
                raise ValueError(msg)
            self.data["amplitudes"] = np.broadcast_to(0.0, (0,))
        else:
            self.data["amplitudes"] = np.broadcast_to(value, (self.modes,))
        self.check_data()

    @abstractmethod
    def interface_distance(self, *angles: RealArray) -> RealArray:
        pass

    @abstractmethod
    def interface_curvature(self, *angles: RealArray) -> RealArray:  # type: ignore
        pass

    @property
    def volume(self) -> float:
        raise NotImplementedError

    @volume.setter
    def volume(self, volume: float) -> None:
        raise NotImplementedError

    @property
    def surface_area(self) -> float:
        raise NotImplementedError

    def _get_phase_field(self, grid: GridBase, dtype: DTypeLike = float) -> RealArray:
        """Creates a normalized image of the droplet on the `grid`

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                The grid used for discretizing the droplet phase field

        Returns:
            :class:`~numpy.ndarray`: An array with data values representing the droplet
            phase field at support points of the `grid`.
        """
        if self.dim != grid.dim:
            msg = (
                f"Droplet (dimension {self.dim}) incompatible with grid (dimension "
                f"{grid.dim})"
            )
            raise ValueError(msg)

        if self.interface_width is None:
            interface_width = grid.typical_discretization
        else:
            interface_width = self.interface_width

        # calculate grid distance from droplet center
        dist, *angles = spherical.polar_coordinates(
            grid, origin=self.position, ret_angle=True
        )

        # calculate interface distance from droplet center
        interface = self.interface_distance(*angles)

        # make the image
        if interface_width == 0 or np.issubdtype(dtype, bool):
            result = dist < interface
        else:
            result = 0.5 + 0.5 * np.tanh(2 * (interface - dist) / interface_width)

        return result.astype(dtype)

    def _get_mpl_patch(self, dim=None, *, color=None, **kwargs):
        """Return the patch representing the droplet for plotting.

        Args:
            dim (int, optional):
                The dimension in which the data is plotted. If omitted, the actual
                physical dimension is assumed
        """
        msg = f"Plotting {self.__class__.__name__} not implemented"
        raise NotImplementedError(msg)


class PerturbedDroplet2D(PerturbedDropletBase):
    r"""Represents a single 2D droplet with a perturbed shape.

    The shape is described using the distance :math:`R(\phi)` of the interface from the
    `position`, which is a function of the polar angle :math:`\phi`. This function is
    expressed as a truncated series of harmonics:

    .. math::
        R(\phi) = R_0 + R_0\sum_{n=1}^N \left[ \epsilon^{(1)}_n \sin(n \phi)
                                             + \epsilon^{(2)}_n \cos(n \phi) \right]

    where :math:`N` is the number of perturbation modes considered, which is given by
    half the length of the `amplitudes` array. Consequently, amplitudes should always
    be an even number, to consider both `sin` and `cos` terms.
    """

    dim = 2

    __slots__ = ["data"]

    def __init__(
        self,
        position: RealArray,
        radius: float,
        interface_width: float | None = None,
        amplitudes: RealArray | None = None,
    ):
        r"""
        Args:
            position (:class:`~numpy.ndarray`):
                Position of the droplet center
            radius (float):
                Radius of the droplet
            interface_width (float, optional):
                Width of the interface
            amplitudes (:class:`~numpy.ndarray`):
                (dimensionless) perturbation amplitudes
                :math:`\{\epsilon^{(1)}_1, \epsilon^{(2)}_1, \epsilon^{(1)}_2,
                \epsilon^{(2)}_2, \epsilon^{(1)}_3, \epsilon^{(2)}_3, \dots \}`.
                The length of the array needs to be even to capture perturbations of the
                highest mode consistently.
        """
        super().__init__(position, radius, interface_width, amplitudes)
        if len(self.amplitudes) % 2 != 0:
            _logger.warning(
                "`amplitudes` should be of even length to capture all perturbations of "
                "the highest mode."
            )

    @enable_scalar_args
    def interface_distance(self, φ: RealArray) -> RealArray:  # type: ignore
        """Calculates the distance of the droplet interface to the origin.

        Args:
            φ (float or :class:`~RealArray`):
                The angle in the polar coordinate system that describing the interface

        Returns:
            Array with distances of the interfacial points associated with each angle φ
        """
        dist: RealArray = np.ones(φ.shape, dtype=float)
        for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
            if a != 0:
                dist += a * np.sin(n * φ)  # type: ignore
            if b != 0:
                dist += b * np.cos(n * φ)  # type: ignore
        return self.radius * dist

    @enable_scalar_args
    def interface_position(self, φ: RealArray) -> RealArray:
        """Calculates the position of the interface of the droplet.

        Args:
            φ (float or :class:`~RealArray`):
                The angle in the polar coordinate system that describing the interface

        Returns:
            Array with coordinates of interfacial points associated with each angle φ
        """
        dist = self.interface_distance(φ)
        pos = dist[:, None] * np.transpose([np.cos(φ), np.sin(φ)])
        return self.position[None, :] + pos  # type: ignore

    @enable_scalar_args
    def interface_curvature(self, φ: RealArray) -> RealArray:  # type: ignore
        r"""Calculates the mean curvature of the interface of the droplet.

        For simplicity, the effect of the perturbations are only included to
        linear order in the perturbation amplitudes :math:`\epsilon^{(1/2)}_n`.

        Args:
            φ (float or :class:`~RealArray`):
                The angle in the polar coordinate system that describing the interface

        Returns:
            Array with curvature at the interfacial points associated with each angle φ
        """
        curv_radius: RealArray = np.ones(φ.shape, dtype=float)
        for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
            factor = n * n - 1
            if a != 0:
                curv_radius -= a * factor * np.sin(n * φ)  # type: ignore
            if b != 0:
                curv_radius -= b * factor * np.cos(n * φ)  # type: ignore
        return 1 / (self.radius * curv_radius)

    @property
    def volume(self) -> float:
        """float: volume of the droplet"""
        term = 1 + np.sum(self.amplitudes**2) / 2
        return np.pi * self.radius**2 * term

    @volume.setter
    def volume(self, volume: float) -> None:
        """Set volume keeping relative perturbations."""
        term = 1 + np.sum(self.amplitudes**2) / 2
        self.radius = np.sqrt(volume / (np.pi * term))

    @property
    def surface_area(self) -> float:
        """float: surface area of the droplet"""
        # discretize surface for simple approximation to integral
        φs, dφ = np.linspace(0, 2 * np.pi, 256, endpoint=False, retstep=True)

        dist: RealArray = np.ones(φs.shape, dtype=float)
        dist_dφ: RealArray = np.zeros(φs.shape, dtype=float)
        for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
            if a != 0:
                dist += a * np.sin(n * φs)  # type: ignore
                dist_dφ += a * n * np.cos(n * φs)  # type: ignore
            if b != 0:
                dist += b * np.cos(n * φs)  # type: ignore
                dist_dφ -= b * n * np.sin(n * φs)  # type: ignore

        dx = dist_dφ * np.cos(φs) - dist * np.sin(φs)
        dy = dist_dφ * np.sin(φs) + dist * np.cos(φs)
        line_element = np.hypot(dx, dy)

        return self.radius * line_element.sum() * dφ  # type: ignore

    @property
    def surface_area_approx(self) -> float:
        """float: surface area of the droplet (quadratic in amplitudes)"""
        length: float = 4
        for n, (a, b) in enumerate(iterate_in_pairs(self.amplitudes), 1):  # no 0th mode
            length += n**2 * (a**2 + b**2)
        return np.pi * self.radius * length / 2

    def _get_mpl_patch(self, dim=2, **kwargs):
        """Return the patch representing the droplet for plotting.

        Args:
            dim (int, optional):
                The dimension in which the data is plotted. If omitted, the actual
                physical dimension is assumed
            **kwargs:
                Additional keyword arguments are passed to
                :class:`matplotlib.patches.Polygon`, which creates the patch that
                represents the droplet. For instance, to only draw the outlines of the
                droplets, you may need to supply `fill=False`.

        Returns:
            :class:`~matplotlib.patches.Polygon`: The patch representing the droplet
        """
        import matplotlib as mpl

        if dim is None:
            dim = self.dim

        if dim != 2:
            msg = "Plotting is only implemented in 2d"
            raise NotImplementedError(msg)

        φ = np.linspace(0, 2 * np.pi, endpoint=False)
        xy = self.interface_position(φ)
        return mpl.patches.Polygon(xy, closed=True, **kwargs)


class PerturbedDroplet3D(PerturbedDropletBase):
    r"""Represents a single 3D droplet with a perturbed shape.

    The shape is described using the distance :math:`R(\theta, \phi)` of the interface
    from the origin as a function of the azimuthal angle :math:`\theta` and the polar
    angle :math:`\phi`. This function is developed as a truncated series of spherical
    harmonics :math:`Y_{l,m}(\theta, \phi)`:

    .. math::
        R(\theta, \phi) = R_0 \left[1 + \sum_{l=1}^{N_l}\sum_{m=-l}^l
                                            \epsilon_{l,m} Y_{l,m}(\theta, \phi) \right]

    where :math:`N_l` is the number of perturbation modes considered, which is deduced
    from the length of the `amplitudes` array.
    """

    dim = 3

    __slots__ = ["data"]

    def __init__(
        self,
        position: RealArray,
        radius: float,
        interface_width: float | None = None,
        amplitudes: RealArray | None = None,
    ):
        r"""
        Args:
            position (:class:`~numpy.ndarray`):
                Position of the droplet center
            radius (float):
                Radius of the droplet
            interface_width (float, optional):
                Width of the interface
            amplitudes (:class:`~numpy.ndarray`):
                Perturbation amplitudes :math:`\epsilon_{l,m}`. Note that the zero-th
                mode, which would only change the radius, is skipped. Consequently, the
                length of the array needs to be 0, 3, 8, 15, 24, ... to capture
                perturbations of the highest mode consistently.
        """
        super().__init__(position, radius, interface_width, amplitudes)
        num_modes = len(self.amplitudes) + 1
        if not spherical.spherical_index_count_optimal(num_modes):
            l, _ = spherical.spherical_index_lm(num_modes)
            opt_modes = spherical.spherical_index_count(l) - 1
            _logger.warning(
                "The length of `amplitudes` should be such that all orders are "
                "captured for the perturbations with the highest degree (%d). "
                "Consider increasing the size of the array to %d.",
                l,
                opt_modes,
            )

    @enable_scalar_args
    def interface_distance(self, θ: RealArray, φ: RealArray | None = None) -> RealArray:  # type: ignore
        r"""Calculates the distance of the droplet interface to the origin.

        Args:
            θ (float or :class:`~RealArray`):
                Azimuthal angle (in :math:`[0, \pi]`)
            φ (float or :class:`~RealArray`):
                Polar angle (in :math:`[0, 2\pi]`); 0 if omitted

        Returns:
            Array with distances of the interfacial points associated with the angles
        """
        if φ is None:
            φ = np.zeros_like(θ)
        elif θ.shape != φ.shape:
            msg = "Shape of θ and φ must agree"
            raise ValueError(msg)
        dist: RealArray = np.ones(θ.shape, dtype=float)
        for k, a in enumerate(self.amplitudes, 1):  # skip zero-th mode!
            if a != 0:
                dist += a * spherical.spherical_harmonic_real_k(k, θ, φ)  # type: ignore
        return self.radius * dist

    @enable_scalar_args
    def interface_position(self, θ: RealArray, φ: RealArray | None = None) -> RealArray:
        r"""Calculates the position of the interface of the droplet.

        Args:
            θ (float or :class:`~RealArray`):
                Azimuthal angle (in :math:`[0, \pi]`)
            φ (float or :class:`~RealArray`):
                Polar angle (in :math:`[0, 2\pi]`); 0 if omitted

        Returns:
            Array with coordinates of the interfacial points associated with the angles
        """
        if φ is None:
            φ = np.zeros_like(θ)
        elif θ.shape != φ.shape:
            msg = "Shape of θ and φ must agree"
            raise ValueError(msg)
        dist = self.interface_distance(θ, φ)
        unit_vector = [np.sin(θ) * np.cos(φ), np.sin(θ) * np.sin(φ), np.cos(θ)]
        pos = dist[:, None] * np.transpose(unit_vector)
        return self.position[None, :] + pos  # type: ignore

    @enable_scalar_args
    def interface_curvature(  # type: ignore
        self, θ: RealArray, φ: RealArray | None = None
    ) -> RealArray:
        r"""Calculates the mean curvature of the interface of the droplet.

        For simplicity, the effect of the perturbations are only included to
        linear order in the perturbation amplitudes :math:`\epsilon_{l,m}`.

        Args:
            θ (float or :class:`~RealArray`):
                Azimuthal angle (in :math:`[0, \pi]`)
            φ (float or :class:`~RealArray`):
                Polar angle (in :math:`[0, 2\pi]`); 0 if omitted

        Returns:
            Array with curvature at the interfacial points associated with the angles
        """
        if φ is None:
            φ = np.zeros_like(θ)
        elif θ.shape != φ.shape:
            msg = "Shape of θ and φ must agree"
            raise ValueError(msg)
        Yk = spherical.spherical_harmonic_real_k
        correction = 0
        for k, a in enumerate(self.amplitudes, 1):  # skip zero-th mode!
            if a != 0:
                l, _ = spherical.spherical_index_lm(k)
                hk = (l**2 + l - 2) / 2
                correction = a * hk * Yk(k, θ, φ)  # type: ignore
        return 1 / self.radius + correction / self.radius**2  # type: ignore

    @property
    def volume(self) -> float:
        """float: volume of the droplet (determined numerically)"""

        def integrand(θ, φ):
            """Helper function calculating the integrand."""
            r = self.interface_distance(θ, φ)
            return r**3 * np.sin(θ) / 3

        volume = integrate.dblquad(
            integrand, 0, 2 * np.pi, lambda _: 0, lambda _: np.pi
        )[0]
        return volume  # type: ignore

    @volume.setter
    def volume(self, volume: float) -> None:
        """Set volume keeping relative perturbations."""
        msg = "Cannot set volume"
        raise NotImplementedError(msg)

    @property
    def volume_approx(self) -> float:
        """float: approximate volume to linear order in the perturbation"""
        volume: float = spherical.volume_from_radius(self.radius, 3)
        if len(self.amplitudes) > 0:
            volume += self.amplitudes[0] * 2 * np.sqrt(np.pi) * self.radius**2
        return volume


class PerturbedDroplet3DAxisSym(PerturbedDropletBase):
    r"""Represents a 3D droplet with axisymmetrically perturbed shape.

    The shape is described using the distance :math:`R(\theta)` of the interface from
    the origin as a function of the azimuthal angle :math:`\theta`, while polar symmetry
    is assumed. This function is developed as a truncated series of spherical harmonics
    :math:`Y_{l,m}(\theta, 0)`:

    .. math::
        R(\theta) = R_0 \left[1 + \sum_{l=1}^{N_l}
                                                \epsilon_{l} Y_{l,0}(\theta, 0) \right]

    where :math:`N_l` is the number of perturbation modes considered, which is deduced
    from the length of the `amplitudes` array.
    """

    dim = 3

    __slots__ = ["data"]

    def check_data(self):
        """Method that checks the validity and consistency of self.data."""
        super().check_data()
        if not np.allclose(self.position[:2], 0):
            msg = "Droplet must lie on z-axis"
            raise ValueError(msg)

    @enable_scalar_args
    def interface_distance(self, θ: RealArray) -> RealArray:  # type: ignore
        r"""Calculates the distance of the droplet interface to the origin.

        Args:
            θ (float or :class:`~RealArray`):
                Azimuthal angle (in :math:`[0, \pi]`)

        Returns:
            Array with distances of the interfacial points associated with the angles
        """
        dist: RealArray = np.ones(θ.shape, dtype=float)
        for order, a in enumerate(self.amplitudes, 1):  # skip zero-th mode!
            if a != 0:
                dist += a * spherical.spherical_harmonic_symmetric(order, θ)  # type: ignore
        return self.radius * dist

    @enable_scalar_args
    def interface_curvature(self, θ: RealArray) -> RealArray:  # type: ignore
        r"""Calculates the mean curvature of the interface of the droplet.

        For simplicity, the effect of the perturbations are only included to
        linear order in the perturbation amplitudes :math:`\epsilon_{l,m}`.

        Args:
            θ (float or :class:`~RealArray`):
                Azimuthal angle (in :math:`[0, \pi]`)

        Returns:
            Array with curvature at the interfacial points associated with the angles
        """
        Yl = spherical.spherical_harmonic_symmetric
        correction = 0
        for order, a in enumerate(self.amplitudes, 1):  # skip zero-th mode!
            if a != 0:
                hl = (order**2 + order - 2) / 2
                correction = a * hl * Yl(order, θ)  # type: ignore
        return 1 / self.radius + correction / self.radius**2  # type: ignore

    @property
    def volume_approx(self) -> float:
        """float: approximate volume to linear order in the perturbation"""
        volume: float = spherical.volume_from_radius(self.radius, 3)
        if len(self.amplitudes) > 0:
            volume += self.amplitudes[0] * 2 * np.sqrt(np.pi) * self.radius**2
        return volume


def droplet_from_data(droplet_class: str, data: RealArray) -> DropletBase:
    """Create a droplet instance of the given class using some data.

    Args:
        droplet_class (str):
            The name of the class that is used to create the droplet instance
        data (:class:`~numpy.ndarray`):
            A numpy array that defines the droplet properties
    """
    cls = DropletBase._subclasses[droplet_class]
    return cls(**{key: data[key] for key in data.dtype.names})  # type: ignore


class _TriangulatedSpheres:
    """Helper class for handling stored data about triangulated spheres."""

    def __init__(self) -> None:
        self.path = Path(__file__).resolve().parent / "resources" / "spheres_3d.hdf5"
        self.num_list = np.zeros((0,))
        self.data: dict[int, dict[str, Any]] | None = None

    def _load(self):
        """Load the stored resource."""
        import h5py

        _logger.info("Open resource `%s`", self.path)
        with h5py.File(self.path, "r") as f:
            self.num_list = np.array(f.attrs["num_list"])
            self.data = {}
            for num in self.num_list:
                group = f[str(num)]
                tri = {
                    "points": np.array(group["points"]),
                    "angles": np.array(group["angles"]),
                    "cells": np.array(group["cells"]),
                }
                self.data[num] = tri

    def get_triangulation(self, num_est: int = 1) -> dict[str, Any]:
        """Get a triangulation of a sphere.

        Args:
            num_est (int): The rough number of vertices in the triangulation

        Returns:
            dict: A dictionary with information about the triangulation
        """
        if self.data is None:
            self._load()

        index = np.argmin((self.num_list - num_est) ** 2)
        return self.data[self.num_list[index]]  # type: ignore


triangulated_spheres = _TriangulatedSpheres()


__all__ = [
    "DiffuseDroplet",
    "PerturbedDroplet2D",
    "PerturbedDroplet3D",
    "PerturbedDroplet3DAxisSym",
    "SphericalDroplet",
]
