"""
Convenience functions that are FiniteElement backend independent.
"""

import numpy as np
from petsc4py import PETSc
from scipy.spatial.distance import cdist
import numpy.typing as npt
from packaging.version import Version

__all__ = ["find_permutation"]

@PETSc.Log.EventDecorator()
def find_permutation(points_a:npt.NDArray[np.inexact], points_b:npt.NDArray[np.inexact],
                     tol:float=1e-5):
    """ Find all permutations between a list of two sets of points.

    Given two numpy arrays of shape (ncells, npoints, dim) containing
    floating point coordinates for each cell, determine each index
    permutation that takes `points_a` to `points_b`. Ie:
    ```
    permutation = find_permutation(points_a, points_b)
    assert np.allclose(points_a[permutation], points_b, rtol=0, atol=tol)
    ```
    """
    if points_a.shape != points_b.shape:
        raise ValueError("`points_a` and `points_b` must have the same shape.")

    p = [np.where(cdist(a, b).T < tol)[1] for a, b in zip(points_a, points_b)]
    try:
        permutation = np.array(p, ndmin=2)
    except ValueError as e:
        raise ValueError(
            "It was not possible to find a permutation for every cell"
            " within the provided tolerance"
        ) from e

    if permutation.shape != points_a.shape[0:2]:
        raise ValueError(
            "It was not possible to find a permutation for every cell"
            " within the provided tolerance"
        )

    return permutation

def trim_util(T):
    """
    Trim zeros from a connectivity array T.
    """
    if Version(np.__version__) >= Version("2.2"):
        T = np.trim_zeros(T, "b", axis=1).astype(np.int32) - 1
    else:
        T = (np.array([list(np.trim_zeros(a, "b")) for a in list(T)],
                dtype=np.int32,
            )- 1)
    return T
