"""
Module for vertex groupings code, especially the GroupingArtist class.
"""

from typing import Union
import numpy as np
import matplotlib as mpl
from matplotlib.collections import PatchCollection


from .typing import (
    GroupingType,
    LayoutType,
)
from .ingest.heuristics import (
    normalise_layout,
    normalise_grouping,
)
from .style import get_style, rotate_style
from .utils.geometry import (
    convex_hull,
    _compute_group_path_with_vertex_padding,
)


class GroupingArtist(PatchCollection):
    """Matplotlib artist for a vertex grouping (clustering/cover).

    This class is used to plot patches surrounding groups of vertices in a network.
    """

    _factor = 1.0

    def __init__(
        self,
        grouping: GroupingType,
        layout: LayoutType,
        vertexpadding: Union[None, int] = None,
        points_per_curve: int = 30,
        transform: mpl.transforms.Transform = mpl.transforms.IdentityTransform(),
        *args,
        **kwargs,
    ) -> None:
        """Container artist for vertex groupings, e.g. covers or clusterings.

        Parameters:
            grouping: This can be a sequence of sets (a la networkx), an igraph Clustering
                or Cover instance (including VertexClustering/VertexCover), or a sequence
                of integers/strings indicating memberships for each vertex.
            layout: The layout of the vertices. If this object has no keys/index, the
                vertices are assumed to have IDs corresponding to integers starting from
                zero.
            vertexpadding: How may points of padding to leave around each vertex centre.
            points_per_curve: How many points to use to approximate a round envelope around
                each convex hull vertex.
            transform: The matplotlib transform to use for the patches (typically transData).
        """
        if vertexpadding is not None:
            self._vertexpadding = vertexpadding
        else:
            style = get_style(".grouping")
            self._vertexpadding = style.get("vertexpadding", 10)

        self._points_per_curve = points_per_curve

        network = kwargs.pop("network", None)
        patches, grouping, coords_hulls = self._create_patches(
            grouping,
            layout,
            network,
            **kwargs,
        )
        if "network" in kwargs:
            del kwargs["network"]
        self._grouping = grouping
        self._coords_hulls = coords_hulls
        kwargs["match_original"] = True

        super().__init__(patches, *args, **kwargs)

        zorder = get_style(".grouping").get("zorder", 1)
        self.set_zorder(zorder)

        self.set_transform(transform)

    def set_figure(self, figure) -> None:
        """Set the figure for the grouping, recomputing the paths depending on the figure's dpi."""
        ret = super().set_figure(figure)
        self._compute_paths(self.get_figure(root=True).dpi)
        return ret

    def get_vertexpadding(self) -> float:
        """Get the vertex padding of each group."""
        return self._vertexpadding

    def get_vertexpadding_dpi(self, dpi: float = 72.0) -> float:
        """Get vertex padding of each group, scaled by dpi of the figure."""
        return self.get_vertexpadding() * dpi / 72.0 * self._factor

    def _create_patches(self, grouping, layout, network, **kwargs):
        layout = normalise_layout(layout, network=network)
        grouping = normalise_grouping(grouping, layout)
        style = get_style(".grouping")
        style.pop("vertexpadding", None)

        style.update(kwargs)

        patches = []
        coords_hulls = []
        for i, (name, vids) in enumerate(grouping.items()):
            if len(vids) == 0:
                continue
            vids = np.array(list(vids))
            coords = layout.loc[vids].values
            idx_hull = convex_hull(coords)
            coords_hull = coords[idx_hull]
            coords_hulls.append(coords_hull)

            stylei = rotate_style(style, i)

            # NOTE: the transform is set later on
            patch = _compute_group_patch_stub(
                coords_hull,
                self._vertexpadding,
                label=name,
                **stylei,
            )

            patches.append(patch)
        return patches, grouping, coords_hulls

    def _compute_paths(self, dpi: float = 72.0) -> None:
        ppc = self._points_per_curve
        for i, hull in enumerate(self._coords_hulls):
            self._paths[i].vertices = _compute_group_path_with_vertex_padding(
                hull,
                self._paths[i].vertices,
                self.get_transform(),
                vertexpadding=self.get_vertexpadding_dpi(dpi),
                points_per_curve=ppc,
            )

    def _process(self) -> None:
        self._compute_paths()

    def draw(self, renderer) -> None:
        """Draw or re-draw the grouping patches.

        Parameters:
            renderer: The renderer to use for drawing the patches.
        """
        # FIXME: this kind of breaks everything since the vertices' magical "_transforms" does
        # not really scale from 72 pixels but rather from the screen's or something.
        # Conclusion: using this keeps consistency across dpis but breaks proportionality of
        # vertexpadding and vertex_size (for now).
        # NOTE: this might be less bad than initially thought in the sense that even perfect
        # scaling does not seem to align the center of the perimeter of the group with the
        # center of the perimeter of the vertex when of the same exact size. So we are
        # probably ok winging it as users will adapt.
        self._compute_paths(self.get_figure(root=True).dpi)
        super().draw(renderer)


def _compute_group_patch_stub(
    points,
    vertexpadding,
    **kwargs,
):
    if vertexpadding == 0:
        return mpl.patches.Polygon(
            points,
            **kwargs,
        )

    # NOTE: Closing point: mpl is a bit quirky here
    vertices = np.zeros(
        (1 + 30 * len(points), 2),
    )
    codes = ["MOVETO"] + ["LINETO"] * (len(vertices) - 2) + ["CLOSEPOLY"]
    codes = [getattr(mpl.path.Path, x) for x in codes]
    patch = mpl.patches.PathPatch(
        mpl.path.Path(
            vertices,
            codes=codes,
        ),
        **kwargs,
    )

    return patch
