# Copyright (C) 2023 Jørgen Schartum Dokken
#
# This file is part of io4dolfinx
#
# SPDX-License-Identifier:    MIT

"""
Vectorized numpy operations used internally in io4dolfinx
"""

from __future__ import annotations

from pathlib import Path

from mpi4py import MPI

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

__all__ = [
    "check_file_exists",
    "compute_local_range",
    "index_owner",
    "compute_dofmap_pos",
    "unroll_dofmap",
    "compute_insert_position",
    "unroll_insert_position",
    "reconstruct_mesh",
]


def check_file_exists(filename: Path | str):
    """Check if file exists."""
    if not Path(filename).exists():
        raise FileNotFoundError(f"{filename} not found")


valid_function_types = np.float32 | np.float64 | np.complex64 | np.complex128
valid_real_types = np.float32 | np.float64


def element_signature(V):
    if Version(dolfinx.__version__) > Version("0.9.0"):
        return V.element.signature
    else:
        return V.element.signature()


def compute_insert_position(
    data_owner: npt.NDArray[np.int32],
    destination_ranks: npt.NDArray[np.int32],
    out_size: npt.NDArray[np.int32],
) -> npt.NDArray[np.int32]:
    """
    Giving a list of ranks, compute the local insert position for each rank in a list
    sorted by destination ranks. This function is used for packing data from a
    given process to its destination processes.

    Example:

        .. highlight:: python
        .. code-block:: python

            data_owner = [0, 1, 1, 0, 2, 3]
            destination_ranks = [2,0,3,1]
            out_size = [1, 2, 1, 2]
            insert_position = compute_insert_position(data_owner, destination_ranks, out_size)

        Insert position is then ``[1, 4, 5, 2, 0, 3]``
    """
    process_pos_indicator = data_owner.reshape(-1, 1) == destination_ranks

    # Compute offsets for insertion based on input size
    send_offsets = np.zeros(len(out_size) + 1, dtype=np.intc)
    send_offsets[1:] = np.cumsum(out_size)
    assert send_offsets[-1] == len(data_owner)

    # Compute local insert index on each process
    proc_row, proc_col = np.nonzero(process_pos_indicator)
    cum_pos = np.cumsum(process_pos_indicator, axis=0)
    insert_position = cum_pos[proc_row, proc_col] - 1

    # Add process offset for each local index
    insert_position += send_offsets[proc_col]
    return insert_position


def unroll_insert_position(
    insert_position: npt.NDArray[np.int32], block_size: int
) -> npt.NDArray[np.int32]:
    """
    Unroll insert position by a block size

    Example:


        .. highlight:: python
        .. code-block:: python

            insert_position = [1, 4, 5, 2, 0, 3]
            unrolled_ip = unroll_insert_position(insert_position, 3)

        where ``unrolled_ip = [3, 4 ,5, 12, 13, 14, 15, 16, 17, 6, 7, 8, 0, 1, 2, 9, 10, 11]``
    """
    unrolled_ip = np.repeat(insert_position, block_size) * block_size
    unrolled_ip += np.tile(np.arange(block_size), len(insert_position))
    return unrolled_ip


def compute_local_range(comm: MPI.Intracomm, N: np.int64):
    """
    Divide a set of `N` objects into `M` partitions, where `M` is
    the size of the MPI communicator `comm`.

    NOTE: If N is not divisible by the number of ranks, the first `r`
    processes gets an extra value

    Returns the local range of values
    """
    rank = comm.rank
    size = comm.size
    n = N // size
    r = N % size
    # First r processes has one extra value
    if rank < r:
        return [rank * (n + 1), (rank + 1) * (n + 1)]
    else:
        return [rank * n + r, (rank + 1) * n + r]


def index_owner(
    comm: MPI.Intracomm, indices: npt.NDArray[np.int64], N: np.int64
) -> npt.NDArray[np.int32]:
    """
    Find which rank (local to comm) which owns an `index`, given that
    data of size `N` has been split equally among the ranks.

    NOTE: If `N` is not divisible by the number of ranks, the first `r`
    processes gets an extra value.
    """
    size = comm.size
    assert (indices < N).all()
    n = N // size
    r = N % size

    owner = np.empty_like(indices, dtype=np.int32)
    inc_remainder = indices < (n + 1) * r
    owner[inc_remainder] = indices[inc_remainder] // (n + 1)
    owner[~inc_remainder] = r + (indices[~inc_remainder] - r * (n + 1)) // n
    return owner


def unroll_dofmap(dofs: npt.NDArray[np.int32], bs: int) -> npt.NDArray[np.int32]:
    """
    Given a two-dimensional dofmap of size `(num_cells, num_dofs_per_cell)`
    Expand the dofmap by its block size such that the resulting array
    is of size `(num_cells, bs*num_dofs_per_cell)`
    """
    num_cells, num_dofs_per_cell = dofs.shape
    unrolled_dofmap = np.repeat(dofs, bs).reshape(num_cells, num_dofs_per_cell * bs) * bs
    unrolled_dofmap += np.tile(np.arange(bs), num_dofs_per_cell)
    return unrolled_dofmap


def compute_dofmap_pos(
    V: dolfinx.fem.FunctionSpace,
) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
    """
    Compute a map from each owned dof in the dofmap to a single cell owned by the
    process, and the relative position of the dof.

    :param V: The function space
    :returns: The tuple (`cells`, `dof_pos`) where each array is the size of the
        number of owned dofs (unrolled for block size)
    """
    dofs = V.dofmap.list
    mesh = V.mesh
    num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
    dofmap_bs = V.dofmap.bs
    num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs

    local_cell = np.empty(
        num_owned_dofs, dtype=np.int32
    )  # Local cell index for each dof owned by process
    dof_pos = np.empty(num_owned_dofs, dtype=np.int32)  # Position in dofmap for said dof

    unrolled_dofmap = unroll_dofmap(dofs[:num_owned_cells, :], dofmap_bs)
    markers = unrolled_dofmap < num_owned_dofs
    local_indices = np.broadcast_to(np.arange(markers.shape[1]), markers.shape)
    cell_indicator = np.broadcast_to(
        np.arange(num_owned_cells, dtype=np.int32).reshape(-1, 1),
        (num_owned_cells, markers.shape[1]),
    )
    indicator = unrolled_dofmap[markers].reshape(-1)
    local_cell[indicator] = cell_indicator[markers].reshape(-1)
    dof_pos[indicator] = local_indices[markers].reshape(-1)
    return local_cell, dof_pos


def reconstruct_mesh(mesh: dolfinx.mesh.Mesh, coordinate_element_degree: int) -> dolfinx.mesh.Mesh:
    """
    Make a copy of a mesh and potentially change the element of the coordinate element.

    Note:
        The topology is shared with the original mesh but the geometry is reconstructed.

    Args:
        mesh: Mesh to reconstruct
        coordinate_element_degree: Degree to use for coordinate element

    Returns:
        The new mesh

    """
    # Extract cell properties
    ud = mesh.ufl_domain()
    assert ud is not None
    c_el = ud.ufl_coordinate_element()
    family = c_el.family_name
    lvar = c_el.lagrange_variant
    ct = c_el.cell_type

    # Create new UFL element
    new_c_el = basix.ufl.element(
        family,
        ct,
        coordinate_element_degree,
        shape=(mesh.geometry.dim,),
        lagrange_variant=lvar,
        dtype=mesh.geometry.x.dtype,
    )

    # Extract new node coordinates
    V_tmp = dolfinx.fem.functionspace(mesh, new_c_el)
    gdim = mesh.geometry.dim
    x = V_tmp.tabulate_dof_coordinates()[:, :gdim]

    # Create new geoemtry
    geom_imap = V_tmp.dofmap.index_map
    geom_dofmap = V_tmp.dofmap.list
    num_nodes_local = geom_imap.size_local + geom_imap.num_ghosts
    original_input_indices = geom_imap.local_to_global(np.arange(num_nodes_local, dtype=np.int32))
    coordinate_element = dolfinx.fem.coordinate_element(
        mesh.topology.cell_type, coordinate_element_degree, lvar, dtype=mesh.geometry.x.dtype
    )
    # Could use create_geometry here when things are fixed
    geom = dolfinx.mesh.Geometry(
        type(mesh.geometry._cpp_object)(
            geom_imap, geom_dofmap, coordinate_element._cpp_object, x, original_input_indices
        )
    )

    # Create new mesh
    new_top = mesh.topology
    cpp_mesh = type(mesh._cpp_object)(mesh.comm, new_top._cpp_object, geom._cpp_object)
    return dolfinx.mesh.Mesh(cpp_mesh, ufl.Mesh(new_c_el))
