from packaging.version import Version
import dolfinx
import ufl
import numpy as np
import numpy.typing as npt
from .utils import unroll_dofmap

__all__ = ["interpolation_matrix", "prepare_interpolation_data", "interpolate_to_surface_submesh"]

if dolfinx.has_petsc4py:
    from petsc4py import PETSc

    __all__.append("petsc_interpolation_matrix")


def prepare_interpolation_data(
    expr: ufl.core.expr.Expr,
    Q: dolfinx.fem.FunctionSpace,
    interpolation_entities: npt.NDArray[np.int32] | None = None,
) -> npt.NDArray[np.inexact]:
    """Convenience function for preparing data required for assembling the interpolation matrix

    .. math::
        \\begin{align*}
        \\Lambda: V &\\rightarrow Q \\\\
        \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j
        \\end{align*}

    where :math:`l_j` is the dual basis of the space :math:`Q` with basis functions :math:`\\phi_j`,
    and :math:`\\psi_j` are the basis functions of the space :math:`V`.

    Args:
        expr: The UFL expression containing a trial function from space `V`
        Q: Output interpolation space
        interpolation_entities: Entities of the domain of the input space `V` that one
            should evaluate the `expr` at. If not provided, it is assumed that
            we are integrating over all cells in `V` and that `Q` is defined on the same grid.
    Returns:
        Interpolation data per cell, as an numpy array.
    """
    if np.issubdtype(dolfinx.default_scalar_type, np.complexfloating):
        raise NotImplementedError("No complex support")

    # Extract argument from expr (in V)
    arguments = ufl.algorithms.extract_arguments(expr)
    assert len(arguments) == 1
    V = arguments[0].ufl_function_space()

    mesh = V.mesh

    if Q.mesh.topology.dim == V.mesh.topology.dim:
        if interpolation_entities is None:
            tdim = mesh.topology.dim
            num_cells = mesh.topology.index_map(tdim).size_local
            interpolation_entities = np.arange(num_cells, dtype=np.int32)
        else:
            if (ndim := interpolation_entities.ndim) != 1:
                raise ValueError(
                    f"Interpolation entities has wrong input shape, should be 1D, got {ndim}"
                )
            num_cells = len(interpolation_entities)
    elif Q.mesh.topology.dim == V.mesh.topology.dim - 1:
        if interpolation_entities is None:
            raise ValueError(
                "For integration onto a submesh of codim 1,"
                + "the integration entities has to be provided"
            )
        else:
            if (ndim := interpolation_entities.ndim) != 2:
                raise ValueError(
                    f"Interpolation entities has wrong input shape, should be 2D, got {ndim}"
                )
            num_cells = interpolation_entities.shape[0]
    else:
        raise RuntimeError("Only codim-1 interpolation matrices can be defined")

    # Extract quadrature points for expression (in Q space)
    try:
        q_points = Q.element.interpolation_points()
    except TypeError:
        q_points = Q.element.interpolation_points

    # Compile expression
    num_points = q_points.shape[0]
    compiled_expr = dolfinx.fem.Expression(expr, q_points)

    # (num_cells, num_points, num_dofs*bs, expr_value_size)
    array_evaluated = compiled_expr.eval(mesh, interpolation_entities)
    assert np.prod(Q.value_shape) == np.prod(expr.ufl_shape)

    # Get data as (num_cells*num_points,1, expr_shape, num_test_basis_functions*test_block_size)
    expr_size = int(np.prod(expr.ufl_shape))
    array_evaluated = array_evaluated.reshape(
        num_cells * q_points.shape[0], 1, expr_size, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
    )

    # Check if we are dealing with a quadrature element or not.
    # They do not have a complete DOLFINx API, which makes them tricky to use.
    try:
        basix_el = Q.element.basix_element
        Q_vs = basix_el.value_size
        pull_back = basix_el.pull_back
        im = basix_el.interpolation_matrix
    except RuntimeError:
        Q_vs = 1  # If we do not have a basix element, assume value size is 1
        assert isinstance(Q.ufl_element().pullback, ufl.pullback.IdentityPullback)
        pull_back = lambda x: None
        assert Q.element.interpolation_ident
        im = None

    new_array = np.zeros(
        (num_cells * num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs),
        dtype=np.float64,
    )

    # Check if pullback is identity, then we can skip this step
    if not isinstance(Q.ufl_element().pullback, ufl.pullback.IdentityPullback):
        jacobian = dolfinx.fem.Expression(ufl.Jacobian(mesh), q_points)
        detJ = dolfinx.fem.Expression(ufl.JacobianDeterminant(mesh), q_points)
        K = dolfinx.fem.Expression(ufl.JacobianInverse(mesh), q_points)
        jacs = jacobian.eval(mesh, np.arange(num_cells, dtype=np.int32)).reshape(
            num_cells * num_points, mesh.geometry.dim, mesh.topology.dim
        )
        detJs = detJ.eval(mesh, np.arange(num_cells, dtype=np.int32)).flatten()
        Ks = K.eval(mesh, np.arange(num_cells, dtype=np.int32)).reshape(
            num_cells * num_points, mesh.geometry.dim, mesh.topology.dim
        )

        for i in range(V.dofmap.bs * V.dofmap.dof_layout.num_dofs):
            for q in range(Q.dofmap.bs):
                new_array[:, q * Q_vs : (q + 1) * Q_vs, i] = pull_back(
                    array_evaluated[:, :, q * Q_vs : (q + 1) * Q_vs, i], jacs, detJs, Ks
                ).reshape(num_cells * num_points, Q_vs)
        new_array = new_array.reshape(
            num_cells, num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
        )
    else:
        new_array = array_evaluated.reshape(
            num_cells, num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
        )

    interpolated_matrix = np.zeros(
        (
            num_cells,
            Q.dofmap.dof_layout.num_dofs * Q.dofmap.bs,
            V.dofmap.bs * V.dofmap.dof_layout.num_dofs,
        ),
        dtype=np.float64,
    )
    # Check if interpolation matrix of dual operator is identity, then we can use a vectorized
    # version of this step
    if Q.element.interpolation_ident:
        # Smart vectorized version with identity mapping
        if Q.dofmap.bs == 1:
            interpolated_matrix = new_array.transpose(0, 2, 1, 3).reshape(
                new_array.shape[0], new_array.shape[1] * new_array.shape[2], new_array.shape[3]
            )
        else:
            i_scalar = new_array.transpose(0, 2, 1, 3)
            interpolated_matrix = np.zeros(
                (new_array.shape[0], new_array.shape[1] * new_array.shape[2], new_array.shape[3])
            )
            for q in range(Q.dofmap.bs):
                interpolated_matrix[:, q :: Q.dofmap.bs, :] = i_scalar[:, q, :, :]

    else:
        # Tedious non-identity version
        for c in range(num_cells):
            for i in range(V.dofmap.bs * V.dofmap.dof_layout.num_dofs):
                tmp_array = np.zeros((int(num_points), Q.dofmap.bs * Q_vs), dtype=np.float64)
                for p in range(num_points):
                    tmp_array[p] = new_array[c, p, :, i]
                if Q.dofmap.bs == 1:
                    interpolated_matrix[c, :, i] = (im @ tmp_array.T.flatten()).flatten()
                else:
                    for q in range(Q.dofmap.bs):
                        interpolated_matrix[c, q :: Q.dofmap.bs, i] = (
                            im @ tmp_array.T[q].flatten()
                        ).flatten()
    # Apply dof transformation to each column (using Piopla maps)
    mesh.topology.create_entity_permutations()
    if Q.element.needs_dof_transformations:
        cell_perm = mesh.topology.get_cell_permutation_info()[:num_cells]

        permuted_matrix = interpolated_matrix.flatten().copy()
        Q.element.Tt_inv_apply(
            permuted_matrix, cell_perm, V.dofmap.bs * V.dofmap.dof_layout.num_dofs
        )
    else:
        permuted_matrix = interpolated_matrix.flatten()
    return permuted_matrix.reshape(interpolated_matrix.shape)


def interpolation_matrix(
    expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace
) -> dolfinx.la.MatrixCSR:
    """Create the interpolation matrix :math:`\\Lambda` of a
    :py:class:`UFL-expression<ufl.core.expr.Expr>` such that

    .. math::
        \\begin{align*}
        \\Lambda: V &\\rightarrow Q \\\\
        \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j
        \\end{align*}

    where :math:`l_j` is the dual basis of the space :math:`Q` with
    basis functions :math:`\\phi_j`, and :math:`\\psi_j` are the basis functions of the
    space :math:`V`.

    Args:
        expr: The UFL expression
        Q: Output interpolation space

    Returns:
        Interpolation matrix as a :py:class:`MatrixCSR<dolfinx.la.MatrixCSR>`.
    """

    arguments = ufl.algorithms.extract_arguments(expr)
    assert len(arguments) == 1
    V = arguments[0].ufl_function_space()

    interpolation_data = prepare_interpolation_data(expr, Q)

    q = ufl.TestFunction(Q)
    a = dolfinx.fem.form(ufl.inner(expr, q) * ufl.dx)

    def scatter(
        A: dolfinx.la.MatrixCSR,
        num_cells: int,
        dofs_visited: npt.NDArray[np.int32],
        num_rows_local: int,
        array_evaluated: npt.NDArray[np.inexact],
        dofmap0: npt.NDArray[np.int32],
        dofmap1: npt.NDArray[np.int32],
    ):
        A.data[:] = 0
        for i in range(num_cells):
            rows = dofmap0[i, :]
            cols = dofmap1[i, :]
            A_local = array_evaluated[i].reshape(len(rows), len(cols))
            row_filter = (dofs_visited[rows] == 1) | (rows >= num_rows_local)
            A_local[row_filter] = 0
            A.add(A_local.flatten(), rows, cols)
            dofs_visited[rows] = 1

    A = dolfinx.fem.create_matrix(a)  # , dolfinx.la.BlockMode.expanded)

    row_dofmap = unroll_dofmap(Q.dofmap.list, Q.dofmap.bs)  # (num_cells, num_rows)
    col_dofmap = unroll_dofmap(V.dofmap.list, V.dofmap.bs)  # (num_cells, num_cols)

    num_cells = Q.mesh.topology.index_map(Q.mesh.topology.dim).size_local
    dofs_visited = np.zeros(
        (Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts) * Q.dofmap.index_map_bs,
        dtype=np.int8,
    )
    num_rows_local = Q.dofmap.index_map.size_local * Q.dofmap.bs
    scatter(A, num_cells, dofs_visited, num_rows_local, interpolation_data, row_dofmap, col_dofmap)
    A.scatter_reverse()
    return A


if dolfinx.has_petsc4py:

    def petsc_interpolation_matrix(
        expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace, use_petsc: bool = False
    ) -> PETSc.Mat:
        """Create the interpolation matrix :math:`\\Lambda` of a
        :py:class:`UFL-expression<ufl.core.expr.Expr>` such that

        .. math::
            \\begin{align*}
            \\Lambda: V &\\rightarrow Q \\\\
            \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j
            \\end{align*}

        where :math:`l_j` is the dual basis of the space :math:`Q` with basis
        functions :math:`\\phi_j`, and :math:`\\psi_j` are the basis functions
        of the space :math:`V`.

        Args:
            expr: The UFL expression
            Q: Output interpolation space

        Returns:
            Interpolation matrix as a :py:class:`PETSc.Mat<petsc4py.PETSc.Mat>`.
        """
        arguments = ufl.algorithms.extract_arguments(expr)
        assert len(arguments) == 1
        V = arguments[0].ufl_function_space()

        interpolation_data = prepare_interpolation_data(expr, Q)

        q = ufl.TestFunction(Q)
        a = dolfinx.fem.form(ufl.inner(expr, q) * ufl.dx)
        A = dolfinx.fem.petsc.create_matrix(a)

        def scatter(
            A: PETSc.Mat,
            num_cells: int,
            dofs_visited: npt.NDArray[np.int32],
            num_rows_local: int,
            array_evaluated: npt.NDArray[np.inexact],
            dofmap0: npt.NDArray[np.int32],
            dofmap1: npt.NDArray[np.int32],
        ):
            A.zeroEntries()
            for i in range(num_cells):
                rows = dofmap0[i, :]
                cols = dofmap1[i, :]
                A_local = array_evaluated[i].reshape(len(rows), len(cols))
                row_filter = (dofs_visited[rows] == 1) | (rows >= num_rows_local)
                A_local[row_filter] = 0
                A.setValuesLocal(rows, cols, A_local, addv=PETSc.InsertMode.ADD_VALUES)
                dofs_visited[rows] = 1

        row_dofmap = unroll_dofmap(Q.dofmap.list, Q.dofmap.bs)  # (num_cells, num_rows)
        col_dofmap = unroll_dofmap(V.dofmap.list, V.dofmap.bs)  # (num_cells, num_cols)
        num_cells = Q.mesh.topology.index_map(Q.mesh.topology.dim).size_local
        dofs_visited = np.zeros(
            (Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts) * Q.dofmap.index_map_bs,
            dtype=np.int8,
        )
        num_rows_local = Q.dofmap.index_map.size_local * Q.dofmap.bs
        scatter(
            A, num_cells, dofs_visited, num_rows_local, interpolation_data, row_dofmap, col_dofmap
        )
        A.assemble()
        return A


def interpolate_to_surface_submesh(
    u_volume: dolfinx.fem.Function,
    u_surface: dolfinx.fem.Function,
    submesh_facets: npt.NDArray[np.int32],
    integration_entities: npt.NDArray[np.int32],
):
    """
    Interpolate a function `u_volume` into the function `u_surface`.
    Note:
        Does not work for DG as no dofs are associated with the facets
    Args:
        u_volume: Function to interpolate data from
        u_surface: Function to interpolate data to
        submesh_facets: Cells in facet mesh
        integration_entities: Integration entities on the parent mesh
            corresponding to the facets in `submesh_facets`
    """
    if Version(dolfinx.__version__) < Version("0.10.0"):
        raise RuntimeError("interpolate_to_submesh requires dolfinx version 0.10.0 or higher")
    V_vol = u_volume.function_space
    mesh = V_vol.mesh
    V_surf = u_surface.function_space
    submesh = V_surf.mesh
    ip = V_surf.element.interpolation_points

    expr = dolfinx.fem.Expression(u_volume, ip)

    mesh.topology.create_connectivity(mesh.topology.dim, submesh.topology.dim)
    mesh.topology.create_connectivity(submesh.topology.dim, mesh.topology.dim)

    data = expr.eval(mesh, integration_entities)

    submesh.topology.create_entity_permutations()
    mesh.topology.create_entity_permutations()
    cell_info = mesh.topology.get_cell_permutation_info()
    ft = V_surf.element.basix_element.cell_type

    for i in range(integration_entities.shape[0]):
        perm = np.arange(data.shape[1], dtype=np.int32)
        V_vol.element.basix_element.permute_subentity_closure_inv(
            perm,
            cell_info[integration_entities[i, 0]],
            ft,
            int(integration_entities[i, 1]),
        )
        data[i] = data[i][perm]

    if len(data.shape) == 3:
        # Data is now (num_cells, value_size,num_points)
        data = data.swapaxes(1, 2)
        # Data is now (value_size, num_cells, num_points)
        data = data.swapaxes(0, 1)

    if expr.value_size == 1:
        shaped_data = data.flatten()
    else:
        shaped_data = data.reshape(expr.value_size, -1)

    if hasattr(u_surface._cpp_object, "interpolate_f"):
        interpolate_func = u_surface._cpp_object.interpolate_f
    else:
        interpolate_func = u_surface._cpp_object.interpolate
    interpolate_func(shaped_data, submesh_facets)
    u_surface.x.scatter_forward()
