import os
import meshio
import trimesh
import warnings
import itertools
import numpy as np
from nibabel.affines import apply_affine
from tqdm import tqdm
import nibabel as nib
from scipy.interpolate import griddata
from vtkmodules.util import numpy_support  # don't import vtk due to opengl issues when running headless
from vtkmodules.util.vtkConstants import VTK_TRIANGLE
from vtkmodules.vtkIOXML import vtkXMLUnstructuredGridWriter
from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid
from vtkmodules.vtkCommonCore import vtkPoints, vtkDoubleArray
from vtkmodules.vtkFiltersCore import vtkCellDataToPointData, vtkPointDataToCellData

import pynibs


def cell_data_to_point_data(tris, data_tris, nodes, method='nearest'):
    """
    A wrapper for scipy.interpolate.griddata to interpolate cell data to node data.

    Parameters
    ----------
    tris : np.ndarray
        (n_tri, 3) element number list.
    data_tris : np.ndarray
        (n_tri x 3) data in tris.
    nodes : np.ndarray
        (n_nodes, 3) nodes coordinates.
    method: str, default: 'nearest'
        Which method to use for interpolation. Default uses NearestNDInterpolator.

    Returns
    -------
    data_nodes : np.ndarray
        Data in nodes.
    """
    elms_center = np.mean(nodes[tris], axis=1)
    return griddata(elms_center, data_tris, nodes, method)


def data_nodes2elements(data, con):
    """
    Transforms data in nodes to elements (triangles or tetrahedra).

    Parameters
    ----------
    data : np.ndarray of float
        (N_nodes, N_data) Data given in the nodes.
    con : np.ndarray of int
        triangles: (N_elements, 3).
        tetrahedra: (N_elements, 4).
        Connectivity index list forming the elements.

    Returns
    -------
    out : np.ndarray of float
        (N_elements, N_data) Data given in the element centers.
    """
    return np.average(data[con], axis=1)


def data_elements2nodes(data, con, precise=False):
    """
    Transforms data in elements (triangles or tetrahedra) to nodes.
    Data can be list of multiple data arrays.

    Parameters
    ----------
    data : np.ndarray of float or list of np.ndarray
        (N_elements, N_data) Data given in the elements (multiple datasets who fit to con may be passed in a list).
    con : np.ndarray of int
        triangles: (N_elements. 3).
        tetrahedra: (N_elements, 4).
        Connectivity index list forming the elements.
    precise : bool, default: False
        Compute data transformation precisely but slow. Better for near-0 values.

    Returns
    -------
    out : np.ndarray of float or list of np.ndarray
        (N_nodes, N_data) Data in nodes.
    """
    # check if single dataset or a list of multiple datasets is passed
    if type(data) is not list:
        single_array_input = True
        data = [data]
    else:
        single_array_input = False

    n_elements = data[0].shape[0]
    n_nodes = con.max() - con.min() + 1
    if con.min() != 0:
        warnings.warn("Node number list is not zero based")

    # built connectivity matrix
    if not precise:
        try:
            c = np.zeros([n_elements, n_nodes])

            for i in range(n_elements):
                c[i, (con[i])] = 1.0 / con.shape[1]

            # filter out nan from dataset
            for i in range(len(data)):
                data[i][np.isnan(data[i])] = 0

            # determine inverse of node matrix
            cinv = np.linalg.pinv(c)
            # transform data from element center to element nodes
            data_in_nodes = [np.dot(cinv, d) for d in data]

            # if single array was provided, return array as well
            if single_array_input:
                data_in_nodes = np.array(data_in_nodes)

            return data_in_nodes

        except np.core._exceptions._ArrayMemoryError:
            warnings.warn("Cannot allocate enough RAM to do fast data->nodes conversion. "
                          "Falling back to (slow) iterative mapping.")
    data_in_nodes = []

    con_flat = con.flatten()
    n_dims = con.shape[1]

    for d in data:
        data_flat = np.repeat(d, n_dims)
        # data_nodes = np.zeros(n_nodes, )
        # for i in tqdm(range(n_nodes), desc="Mapping elements to node data"):
        # data_nodes[i] = d[np.argwhere(con == i)[:, 0]].mean()

        data_in_nodes.append(np.array([data_flat[con_flat == i].mean() for i in range(n_nodes)]))

    # if single array was provided, return array as well
    if single_array_input:
        data_in_nodes = np.array(data_in_nodes)
    return data_in_nodes


def project_on_scalp_hdf5(coords, mesh, scalp_tag=1005):
    """
    Find the node in the scalp closest to each coordinate.

    Parameters
    ----------
    coords: np.ndarray
        (n, 3) Vectors to be transformed.
    mesh: str or pynibs.TetrahedraLinear
        Filename of mesh in .hdf5 format or Mesh structure.
    scalp_tag: int, default: 1005
        Tag in the mesh where the scalp is to be set.

    Returns
    -------
    points_closest: np.ndarray
        (n, 3) coordinates projected scalp (closest skin points).
    """
    # read head mesh and extract skin surface
    if isinstance(mesh, str):
        mesh = pynibs.hdf5_io.load_mesh_hdf5(mesh)

    if coords.ndim == 1:
        coords = coords[np.newaxis,]

    # crop to skin surface
    triangles_skin = mesh.triangles[mesh.triangles_regions == scalp_tag]
    point_idx_skin = np.unique(triangles_skin)
    points_skin = mesh.points[point_idx_skin]

    # find points with smalled Euclidean distance
    points_closest = np.zeros(coords.shape)
    for i, c in enumerate(coords):
        points_closest[i,] = points_skin[np.argmin(np.linalg.norm(points_skin - c, axis=1)),]

    return points_closest


def project_on_scalp(coords, mesh, scalp_tag=1005):
    """
    Find the node in the scalp closest to each coordinate.

    Parameters
    ----------
    coords: nx3 np.ndarray
        Vectors to be transformed.
    mesh: pynibs.TetrahedraLinear or simnibs.msh.mesh_io.Msh
        Mesh structure in simnibs or pynibs format.
    scalp_tag: int, default: 1005
        Tag in the mesh where the scalp is to be set.

    Returns
    -------
    points_closest: np.ndarry
        (n, 3) coordinates projected scalp (closest skin points).
    """
    from simnibs.msh.transformations import project_on_scalp as project_on_scalp_msh
    from simnibs.msh.mesh_io import Msh

    if isinstance(mesh, pynibs.mesh.TetrahedraLinear):
        points_closest = project_on_scalp_hdf5(coords=coords, mesh=mesh, scalp_tag=scalp_tag)
    elif isinstance(mesh, Msh):
        points_closest = project_on_scalp_msh(coords=coords, mesh=mesh, scalp_tag=scalp_tag, distance=0.)
    else:
        raise ValueError(f"Unknown mesh type: {type(mesh)}.")

    return points_closest


def refine_surface(fn_surf, fn_surf_refined, center, radius, repair=True, remesh=True, verbose=True):
    """
    Refines surface (.stl) in spherical ROI and saves as .stl file.

    Parameters
    ----------
    fn_surf : str
        Input filename (.stl).
    fn_surf_refined : str
        Output filename (.stl).
    center : np.ndarray of float
        (3) Center of spherical ROI (x,y,z).
    radius : float
        Radius of ROI.
    repair : bool, default: True
        Repair surface mesh to ensure that it is watertight and forms a volume.
    remesh : bool, default: False
        Perform remeshing with meshfix (also removes possibly overlapping facets and intersections).
    verbose : bool, default: True
        Print output messages.

    Returns
    -------
    <file>: .stl file
    """
    radius_ = radius + 2
    refine = True

    while refine:
        if verbose:
            print(f"Loading {fn_surf} ...")
        # reading original .stl file
        wm = trimesh.load(fn_surf)

        tris = wm.faces
        tris_center = wm.triangles_center
        points = wm.vertices

        # Splitting elements by adding tris_center to points in ROI
        mask_roi = np.linalg.norm(tris_center - center, axis=1) < radius
        ele_idx_roi = np.where(np.linalg.norm(tris_center - center, axis=1) < radius)[0]
        points_refine = points
        tris_refine = tris

        if verbose:
            print(f"Splitting elements ...")

        for ele_idx in tqdm(ele_idx_roi):
            points_idx_ele = tris[ele_idx, :]
            p_0 = points[points_idx_ele[0], :]
            p_1 = points[points_idx_ele[1], :]
            p_2 = points[points_idx_ele[2], :]
            p_01 = p_0 + 0.5 * (p_1 - p_0)
            p_02 = p_0 + 0.5 * (p_2 - p_0)
            p_12 = p_1 + 0.5 * (p_2 - p_1)

            points_refine = np.vstack((points_refine, p_01, p_02, p_12))

            mask_roi = np.hstack((mask_roi, False, False, False, False))

            # add 6 new triangles
            p_0_idx = points_idx_ele[0]
            p_1_idx = points_idx_ele[1]
            p_2_idx = points_idx_ele[2]
            p_01_idx = points_refine.shape[0] - 3
            p_02_idx = points_refine.shape[0] - 2
            p_12_idx = points_refine.shape[0] - 1

            # adding 4 elements
            tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_01_idx, p_02_idx],
                                                            [p_01_idx, p_1_idx, p_12_idx],
                                                            [p_02_idx, p_12_idx, p_2_idx],
                                                            [p_01_idx, p_12_idx, p_02_idx]])))

        ele_idx_del = []

        if radius != np.inf:
            if verbose:
                print(f"Adding triangles in surrounding elements ...")
            # add triangles in surrounding elements
            ele_sur_idx = np.where(np.logical_and(np.linalg.norm(tris_center - center, axis=1) < radius_,
                                                  np.linalg.norm(tris_center - center, axis=1) >= radius))[0]

            for ele_sur in tqdm(ele_sur_idx):
                points_idx_ele = tris[ele_sur, :]
                p_0 = points[points_idx_ele[0], :]
                p_1 = points[points_idx_ele[1], :]
                p_2 = points[points_idx_ele[2], :]
                p_01 = p_0 + 0.5 * (p_1 - p_0)
                p_02 = p_0 + 0.5 * (p_2 - p_0)
                p_12 = p_1 + 0.5 * (p_2 - p_1)

                p_0_idx = points_idx_ele[0]
                p_1_idx = points_idx_ele[1]
                p_2_idx = points_idx_ele[2]

                p_on_02 = False
                p_on_12 = False
                p_on_01 = False

                if (np.sum(p_01 == points_refine, axis=1) == 3).any():
                    p_on_01 = True

                if (np.sum(p_02 == points_refine, axis=1) == 3).any():
                    p_on_02 = True

                if (np.sum(p_12 == points_refine, axis=1) == 3).any():
                    p_on_12 = True

                # no edge with point
                if not p_on_01 and not p_on_02 and not p_on_12:
                    pass

                # one edge with point
                elif p_on_01 and not p_on_02 and not p_on_12:
                    ele_idx_del.append(ele_sur)
                    p_01_idx = np.where(np.sum(points_refine == p_01, axis=1) == 3)[0][0]
                    tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_01_idx, p_2_idx],
                                                                    [p_01_idx, p_1_idx, p_2_idx]])))

                elif p_on_02 and not p_on_01 and not p_on_12:
                    ele_idx_del.append(ele_sur)
                    p_02_idx = np.where(np.sum(points_refine == p_02, axis=1) == 3)[0][0]
                    tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_1_idx, p_02_idx],
                                                                    [p_02_idx, p_1_idx, p_2_idx]])))

                elif p_on_12 and not p_on_02 and not p_on_01:
                    ele_idx_del.append(ele_sur)
                    p_12_idx = np.where(np.sum(points_refine == p_12, axis=1) == 3)[0][0]
                    tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_1_idx, p_12_idx],
                                                                    [p_0_idx, p_12_idx, p_2_idx]])))

                # 2 edges with points
                elif p_on_02 and p_on_12 and not p_on_01:
                    ele_idx_del.append(ele_sur)
                    p_12_idx = np.where(np.sum(points_refine == p_12, axis=1) == 3)[0][0]
                    p_02_idx = np.where(np.sum(points_refine == p_02, axis=1) == 3)[0][0]
                    tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_1_idx, p_02_idx],
                                                                    [p_1_idx, p_12_idx, p_02_idx],
                                                                    [p_02_idx, p_12_idx, p_2_idx]])))

                elif p_on_02 and p_on_01 and not p_on_12:
                    ele_idx_del.append(ele_sur)
                    p_01_idx = np.where(np.sum(points_refine == p_01, axis=1) == 3)[0][0]
                    p_02_idx = np.where(np.sum(points_refine == p_02, axis=1) == 3)[0][0]
                    tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_01_idx, p_02_idx],
                                                                    [p_01_idx, p_2_idx, p_02_idx],
                                                                    [p_01_idx, p_1_idx, p_2_idx]])))

                elif p_on_01 and p_on_12 and not p_on_02:
                    ele_idx_del.append(ele_sur)
                    p_01_idx = np.where(np.sum(points_refine == p_01, axis=1) == 3)[0][0]
                    p_12_idx = np.where(np.sum(points_refine == p_12, axis=1) == 3)[0][0]
                    tris_refine = np.vstack((tris_refine, np.array([[p_0_idx, p_01_idx, p_2_idx],
                                                                    [p_01_idx, p_12_idx, p_2_idx],
                                                                    [p_01_idx, p_1_idx, p_12_idx]])))

        if verbose:
            print("Deleting old triangles ...")

        # delete old triangles
        ele_idx_roi = np.where(mask_roi)[0]
        ele_idx_lst_del = ele_idx_del + list(ele_idx_roi)
        tris_refine = np.delete(tris_refine, ele_idx_lst_del, 0)

        points_refine = np.round_(points_refine, 5)

        # # delete duplicate points
        # p_added = points_refine[points.shape[0]:, :]
        #
        # point_idx_del = np.array([])
        # for i_p, p in tqdm(enumerate(p_added)):
        #
        #     p_idx = np.where(np.sum(p == points_refine, axis=1) == 3)[0]
        #
        #     if len(p_idx) > 1:
        #         if p_idx[1] not in point_idx_del:
        #             point_idx_del = np.hstack((point_idx_del, p_idx[1:]))
        #
        #             # loop over point_idx_del and replace with first point idx
        #             for p_d_idx in p_idx[1:]:
        #                 tris_refine[tris_refine == p_d_idx] = p_idx[0]
        #
        # point_idx_keep = [i for i in range(points_refine.shape[0]) if i not in point_idx_del]
        # point_idx_new = [i for i in range(len(point_idx_keep))]
        # points_refine = points_refine[point_idx_keep, :]
        #
        # # renumber
        # for p_idx_keep, p_idx_new in zip(point_idx_keep[points.shape[0]:], point_idx_new[points.shape[0]:]):
        #     tris_refine[tris_refine == p_idx_keep] = p_idx_new

        # create new trimesh
        mesh = trimesh.Trimesh(vertices=points_refine,
                               faces=tris_refine)

        if repair:
            if mesh.is_watertight:
                if verbose:
                    print(f"Surface is watertight ...")
                mesh_ok = True
            else:
                if verbose:
                    print(f"Surface is NOT watertight ... trying to repair mesh ... ")
                # repair mesh
                trimesh.repair.fill_holes(mesh)

                if mesh.is_watertight:
                    if verbose:
                        print(f"Surface repaired ...")
                    mesh_ok = True

                else:
                    mesh_ok = False
                    radius -= 1
                    radius_ = radius + 2

                    if verbose:
                        print(f"WARNING: Could not repair refined surface ... "
                              f"shrinking radius by 1 mm to {radius} mm")
        else:
            mesh_ok = True

        if mesh_ok:
            if verbose:
                print(f"Saving {fn_surf_refined} ...")
            mesh.export(fn_surf_refined, file_type='stl_ascii')

            if remesh:
                # remesh surface
                print(f"Remeshing {fn_surf_refined} ...")
                command = f"meshfix {fn_surf_refined} -a 2.0 -u 1 -q --shells 9 " \
                          f"--stl -o {fn_surf_refined}"
                os.popen(command).read()

            refine = False


def map_data_to_surface(datasets, points_datasets, con_datasets, fname_fsl_gm, fname_fsl_wm, fname_midlayer=None,
                        delta=0.5, input_data_in_center=True, return_data_in_center=True, data_substitute=-1):
    """
    Maps data from ROI of fsl surface (wm, gm, or midlayer) to given Freesurfer brain surface (wm, gm, inflated).

    Parameters
    ----------
    datasets : np.ndarray of float [N_points x N_data] or list of np.ndarray
        Data in nodes or center of triangles in ROI (specify this in "data_in_center").
    points_datasets : np.ndarray of float [N_points x 3] or list of np.ndarray
        Point coordinates (x,y,z) of ROI where data in datasets list is given, the points have to be a subset of the
        GM/WM surface (has to be provided for each dataset).
    con_datasets : np.ndarray of int [N_tri x 3] or list of np.ndarray
        Connectivity matrix of dataset points (has to be provided for each dataset).
    fname_fsl_gm : str or list of str or list of None
        Filename of pial surface fsl file(s) (one or two hemispheres)
        e.g. in mri2msh: .../fs_ID/surf/lh.pial
    fname_fsl_wm : str or list of str or list of None
        Filename of wm surface fsl file(s) (one or two hemispheres)
        e.g. in mri2msh: .../fs_ID/surf/lh.white
    fname_midlayer : str or list of str
        Filename of midlayer surface fsl file(s) (one or two hemispheres)
        e.g. in headreco: .../fs_ID/surf/lh.central
    delta : float
        Distance parameter where gm-wm surface was generated 0...1 (default: 0.5)
        0 -> WM surface
        1 -> GM surface
    input_data_in_center : bool
        Flag if data in datasets in given in triangle centers or in points (Default: True).
    return_data_in_center : bool
        Flag if data should be returned in nodes or in elements (Default: True).
    data_substitute : float
        Data substitute with this number for all points in the inflated brain,
        which do not belong to the given data set.

    Returns
    -------
    data_mapped : np.ndarray of float [N_points_inf x N_data]
        Mapped data to target brain surface. In points or elements.
    """
    if type(fname_fsl_gm) is not list:
        fname_fsl_gm = [fname_fsl_gm]

    if type(fname_fsl_wm) is not list:
        fname_fsl_wm = [fname_fsl_wm]

    if type(fname_midlayer) is not list:
        fname_midlayer = [fname_midlayer]

    if fname_midlayer[0] is None:
        # load all freesurfer surfaces of gm and wm (hemispheres) and create midlayer
        points_gm = []
        con_target = []
        points_wm = []
        con_idx = 0

        for f_gm, f_wm in zip(fname_fsl_gm, fname_fsl_wm):
            p_gm, c_tar = nib.freesurfer.read_geometry(f_gm)
            p_wm, _ = nib.freesurfer.read_geometry(f_wm)

            points_gm.append(p_gm)
            points_wm.append(p_wm)
            con_target.append(c_tar + con_idx)
            con_idx += np.max(c_tar) + 1  # c_tar.shape[0]

        points_gm = np.vstack(points_gm)
        points_wm = np.vstack(points_wm)
        con_target = np.vstack(con_target)

        # regenerate the gm-wm surface w/o cropping in order to find congruent points
        wm_gm_vector = points_gm - points_wm

        # determine wm-gm surface (midlayer)
        points = points_wm + wm_gm_vector * delta

    else:
        # load directly all freesurfer midlayer surfaces (hemispheres)
        points = []
        con_target = []
        con_idx = 0

        for f_mid in fname_midlayer:
            if f_mid.endswith('.gii'):
                img = nib.load(f_mid)
                p_mid = img.agg_data('pointset')
                c_tar = img.agg_data('triangle')
            else:
                p_mid, c_tar = nib.freesurfer.read_geometry(f_mid)

            points.append(p_mid)
            con_target.append(c_tar + con_idx)
            con_idx += np.max(c_tar) + 1  # c_tar.shape[0]

        points = np.vstack(points)
        con_target = np.vstack(con_target)

    # check datasets
    if type(datasets) is not list:
        datasets = [datasets]

    for i in range(len(datasets)):
        if datasets[i].ndim == 1:
            datasets[i] = datasets[i][:, np.newaxis]
        elif datasets[i].shape[0] < datasets[i].shape[1]:
            raise Warning("Datasets #{} shape[0] dimension is smaller than shape[1] (less points than dataset"
                          "components). Input dimension should be [N_points x N_data] ")

    if type(points_datasets) is not list:
        points_datasets = [points_datasets]

    if type(con_datasets) is not list:
        con_datasets = [con_datasets]
    # check if all points and all con are the same (if so, just map once and reuse results)
    all_points_equal = all([(points_datasets[i] == points_datasets[i + 1]).all()
                            for i in range(len(points_datasets) - 1)])

    all_con_equal = all([(con_datasets[i] == con_datasets[i + 1]).all()
                         for i in range(len(con_datasets) - 1)])

    if all_points_equal and all_con_equal:
        n_main_iter = 1
        n_sub_iter = len(datasets)
    else:
        n_main_iter = len(datasets)
        n_sub_iter = 1

    # check if indexation starts with value greater zero
    if np.min(con_target) > 0:
        con_target = con_target - np.min(con_target)

    n_points = points.shape[0]

    data_mapped = []

    # for i, data in enumerate(datasets):
    for i in range(n_main_iter):
        n_data = datasets[i].shape[1] if datasets[i].ndim > 1 else 1

        # n_points_cropped = points_datasets[i].shape[0]

        # check if indexation starts with value greater zero
        if np.min(con_datasets[i]) > 0:
            con_datasets[i] = con_datasets[i] - np.min(con_datasets[i])

        if datasets[i].ndim == 1:
            datasets[i] = datasets[i][:, np.newaxis]

        if input_data_in_center and return_data_in_center:
            # determine triangle center of dataset
            triangle_center_datasets = np.average(points_datasets[i][con_datasets[i]], axis=1)

            # determine triangle center of whole surface
            triangle_center_surface = np.average(points[con_target], axis=1)

            # loop over all points to get index list
            # point_idx_target = []
            # point_idx_data = []

            point_idx_target = np.zeros(datasets[i].shape[0])
            point_idx_data = np.arange(datasets[i].shape[0])
            for j in tqdm(range(datasets[i].shape[0]), desc="Mapping ROI to surface"):
                point_idx_target[j] = np.where(np.all(np.isclose(triangle_center_datasets[j,], triangle_center_surface),
                                                      axis=1))[0]
            point_idx_target = point_idx_target.astype(int).tolist()
            point_idx_data = point_idx_data.astype(int).tolist()

            # run subiterations (if all points and cons are equal, we save a lot of time here)
            for k in range(n_sub_iter):
                data_mapped.append(np.zeros([triangle_center_surface.shape[0], n_data]) + data_substitute * 1.0)
                data_mapped[k][point_idx_target, :] = datasets[k][point_idx_data, :]

        else:
            # loop over all points to get index list
            point_idx_target = []
            point_idx_data = list(range(datasets[i].shape[0]))

            for j in range(datasets[i].shape[0]):
                point_idx_target.append(np.where(np.all(np.isclose(points_datasets[i][j,], points), axis=1))[0])

            point_idx_target = [int(p) for p in point_idx_target]
            point_idx_data = [int(p) for p in point_idx_data]

            # run subiterations (if all points and cons are equal, we save a lot of time here)
            for k in range(n_sub_iter):
                # transform data from triangle center to triangle nodes if necessary
                if input_data_in_center:
                    data_nodes = data_elements2nodes(datasets[k], con_datasets[k])
                else:
                    data_nodes = datasets[k]

                # find and map data points
                data_mapped.append(np.zeros([n_points, n_data]) + data_substitute * 1.0)
                data_mapped[k][point_idx_target] = data_nodes[point_idx_data]

                # return data in elements instead of points
                if return_data_in_center:
                    data_mapped[k] = data_nodes2elements(data_mapped[k], con_target)

    return data_mapped


def midlayer_2_surf(midlayer_data, coords_target, coords_midlayer, midlayer_con=None, midlayer_data_in_nodes=False,
                    max_dist=5,
                    outside_roi_val=0, precise_map=True):
    """
    Convert midlayer data to whole-brain surface data, e.g. grey matter.
    Output is returned as data in nodes.

    Parameters
    ----------
    midlayer_data : np.ndarray of float
        (n_elm_midlayer,) or (n_nodes_midlayer,), the data in the midlayer.
    coords_target : np.ndarray of float
        (n_nodes_target, 3) Coordinates of the nodes of the target surface.
    coords_midlayer : np.ndarray of float
        (n_nodes_midlayer, 3) Coordinates of the nodes of the midlayer surface.
    midlayer_con : np.ndarray of int, optional
        (n_elm_midlayer, 3) Connectivity of the midlayer elements. Provide if data_in_points == True.
    midlayer_data_in_nodes : bool, default=False
        If midlayer data is provided in nodes, set to True and provide midlayer_con.
    max_dist : float, default=5
        Maximum distance between target and midlayer nodes to pull data from midlayer_data for.
    outside_roi_val : float, default=0
        Areas outside of max_dist are filled with outside_roi_val.
    precise_map : bool, default=True
        If elements to nodes mapping is done, perform this precise and slow or not.

    Returns
    -------
    data_target : np.ndarray
        (n_nodes_target, 1) The data in nodes of the target surface.
    """
    if not midlayer_data_in_nodes:
        assert midlayer_con is not None
        midlayer_data = np.squeeze(pynibs.mesh.data_elements2nodes(midlayer_data, midlayer_con, precise=precise_map))

    data_target = np.zeros((coords_target.shape[0]))
    for i in tqdm(range(data_target.shape[0]), desc='Mapping midlayer2surface'):
        idx = np.linalg.norm(coords_target[i] - coords_midlayer, axis=1).argmin()
        if np.linalg.norm(coords_target[i] - coords_midlayer[idx]) > max_dist:
            data_target[i] = outside_roi_val
        else:
            data_target[i] = midlayer_data[idx]

    return data_target


def point_data_to_cell_data_vtk(mesh=None, nodes=None, con=None, point_data=None, fn=None):
    """
    Convert point data to cell data in a VTK unstructured grid and save the result to a file.

    Parameters
    ----------
    mesh : meshio.Mesh, optional
        The mesh object containing points and cells.
    nodes : np.ndarray of float, optional
        (N_points, 3) Coordinates of the nodes.
    con : np.ndarray of int, optional
        (N_elements, 3) Connectivity index list forming the elements.
    point_data : dict, optional
        Point data to be transformed to cell data.
    fn : str, optional
        If provided, vtk file is written out to this file.

    Returns
    -------
    dict : cell_data
        All data sets from mesh transformed to cell_data.
    """
    if mesh is not None:
        assert nodes is None and con is None, "Provide either mesh or nodes and con."
    else:
        assert nodes is not None and con is not None, "Provide either mesh or nodes and con."
        mesh = meshio.Mesh(points=nodes, cells=[("triangle", con)], point_data=point_data)

    # Create VTK mesh instance from meshio object
    vtk_unstrgrid = vtkUnstructuredGrid()
    points = vtkPoints()
    number_of_points = mesh.points.shape[0]
    number_of_cells = mesh.cells[0].data.shape[0]
    cell_data = {}

    # Insert points into VTK points object
    for idx, p in tqdm(enumerate(mesh.points),
                       desc='Processing points', leave=False):
        points.InsertPoint(idx, p)
    vtk_unstrgrid.SetPoints(points)
    del points

    # Allocate and insert cells into VTK unstructured grid
    vtk_unstrgrid.Allocate(number_of_cells)
    assert len(mesh.cells) == 1, f"Only one cell block is supported, {len(mesh.cells)} present in mesh."
    for idx in tqdm(mesh.cells[0].data,
                    desc='Processing cells', leave=False):
        vtk_unstrgrid.InsertNextCell(VTK_TRIANGLE, 3, idx)

    for arr_name, arr_data_points_meshio in mesh.point_data.items():
        # Create and set point data array
        if len(arr_data_points_meshio.shape) == 1:
            arr_data_points_meshio = arr_data_points_meshio[:, np.newaxis]
        # arr_data_points_meshio = np.atleast_2d(arr_data_points_meshio)
        n_comps = arr_data_points_meshio.shape[1]
        arr_data_points_vtk = vtkDoubleArray()
        arr_data_points_vtk.SetNumberOfComponents(n_comps)
        arr_data_points_vtk.SetNumberOfTuples(number_of_points)
        arr_data_points_vtk.SetName(arr_name)

        for idx, data in tqdm(enumerate(arr_data_points_meshio), total=arr_data_points_meshio.shape[0],
                              desc=f"Processing data {arr_name}", leave=False):
            arr_data_points_vtk.SetTuple(idx, data)
        vtk_unstrgrid.GetPointData().AddArray(arr_data_points_vtk)

        # Convert point data to cell data
        p2c_conv = vtkPointDataToCellData()
        p2c_conv.SetInputData(vtk_unstrgrid)
        p2c_conv.Update()
        ptdata_unstrgrid = p2c_conv.GetOutput()
        arr_data_cell_vtk = ptdata_unstrgrid.GetCellData().GetArray(arr_name)

        # Add cell data array to the unstructured grid
        vtk_unstrgrid.GetCellData().AddArray(arr_data_cell_vtk)
        cell_data[arr_name] = numpy_support.vtk_to_numpy(arr_data_cell_vtk)

    # Write the VTK mesh to a file
    if fn is not None:
        write_vtu(fn, vtk_unstrgrid)

    return cell_data


def cell_data_to_point_data_vtk(mesh=None, nodes=None, con=None, cell_data=None, fn=None):
    """
    Convert cell data to point data in a VTK unstructured grid and save the result to a file.

    Parameters
    ----------
    mesh : meshio.Mesh
        The mesh object containing points and cells.
    nodes : np.ndarray of float, optional
        (N_points, 3) Coordinates of the nodes.
    con : np.ndarray of int, optional
        (N_elements, 3) Connectivity index list forming the elements.
    cell_data : dict, optional
        Cell data to be transformed to point data. keys: str, values: np.ndarray
    fn : str, optional
        If provided, vtk file is written out to this file.

    Returns
    -------
    dict : point_data
        All data sets from mesh transformed to point data.
    """
    if mesh is not None:
        assert nodes is None and con is None, "Provide either mesh or nodes and con."
    else:
        assert nodes is not None and con is not None, "Provide either mesh or nodes and con."
        mesh = meshio.Mesh(points=nodes, cells=[("triangle", con)], cell_data=cell_data)
    # Create VTK mesh instance from meshio object
    vtk_unstrgrid = vtkUnstructuredGrid()
    points = vtkPoints()
    # number_of_points = mesh.points.shape[0]
    number_of_cells = mesh.cells[0].data.shape[0]
    point_data = {}

    # Insert points into VTK points object
    for idx, p in tqdm(enumerate(mesh.points), desc='Processing points', leave=False):
        points.InsertPoint(idx, p)
    vtk_unstrgrid.SetPoints(points)
    del points

    # Allocate and insert cells into VTK unstructured grid
    vtk_unstrgrid.Allocate(number_of_cells)
    assert len(mesh.cells) == 1, f"Only one cell block is supported, {len(mesh.cells)} present in mesh."
    for idx in tqdm(mesh.cells[0].data, desc='Processing cells', leave=False):
        vtk_unstrgrid.InsertNextCell(VTK_TRIANGLE, 3, idx)

    for arr_name, arr_data_cells_meshio in mesh.cell_data.items():
        if isinstance(arr_data_cells_meshio, list):
            if len(arr_data_cells_meshio) > 1:
                raise ValueError
            arr_data_cells_meshio = arr_data_cells_meshio[0]
        # Create and set cell data array
        if len(arr_data_cells_meshio.shape) == 1:
            arr_data_cells_meshio = arr_data_cells_meshio[:, np.newaxis]
        # arr_data_cells_meshio = np.atleast_2d(arr_data_cells_meshio)
        n_comps = arr_data_cells_meshio.shape[1]
        arr_data_cells_vtk = vtkDoubleArray()
        arr_data_cells_vtk.SetNumberOfComponents(n_comps)
        arr_data_cells_vtk.SetNumberOfTuples(number_of_cells)
        arr_data_cells_vtk.SetName(f"{arr_name}")

        for idx, data in tqdm(enumerate(arr_data_cells_meshio), total=arr_data_cells_meshio.shape[1],
                              desc=f"Processing data {arr_name}", leave=False):
            arr_data_cells_vtk.SetTuple(idx, data)
        vtk_unstrgrid.GetCellData().AddArray(arr_data_cells_vtk)

        # Convert cell data to point data
        c2p_conv = vtkCellDataToPointData()
        c2p_conv.SetInputData(vtk_unstrgrid)
        c2p_conv.Update()
        ptdata_unstrgrid = c2p_conv.GetOutput()
        arr_data_point_vtk = ptdata_unstrgrid.GetPointData().GetArray(f"{arr_name}")

        # Add point data array to the unstructured grid
        vtk_unstrgrid.GetPointData().AddArray(arr_data_point_vtk)
        point_data[arr_name] = numpy_support.vtk_to_numpy(arr_data_point_vtk)

    # Write the VTK mesh to a file
    if fn is not None:
        write_vtu(fn, vtk_unstrgrid)

    return point_data


def write_vtu(fn, vtk_grid):
    writer = vtkXMLUnstructuredGridWriter()
    writer.SetFileName(fn)
    writer.SetInputData(vtk_grid)
    writer.Write()


def map_nii2surf(tri_node_coords, nii, datatype=float, nonzeroneighbors=False):
    """
    Map nifti data to mesh surface.

    For each surface triangle,

        - the center coordinate is determined
        - and the corresponding voxel value from the nifti file is extracted.

    Raw, unsmoothed data is returned. Nifti file and mesh needs to be in the same space.

    .. figure:: ../../doc/images/nii_on_surf.png
       :scale: 60 %
       :alt: Nifti volume on mesh surface

       Volume data from nifti file is mapped to mesh surface.

    .. code-block:: python

        fn_roi = "geo.hdf5"
        fn_nii = "some_T_map.nii"

        nii = nib.load(fn_nii)

        with h5py.File(fn_roi, 'r') as f:
            triangle_number_list = f['/mesh/elm/triangle_number_list'][:]
            node_coord = f['/mesh/nodes/node_coord'][:]

        nii_data_on_surf = pynibs.map_nii2surf(node_coord[triangle_number_list], nii)
        output_fn = "some_T_map_on_surf.vtk"
        point_data = {"T": [nii_data_on_surf]}
        meshio.Mesh(np.squeeze(node_coord), [('triangle', triangle_number_list)], cell_data=point_data).write(output_fn)

    Parameters
    ----------
    tri_node_coords: np.ndarray of int
        (n,3,3) list of node coordinates for n trianlges (mesh surface).
    nii: nibabel.Nifti1Image or str
        Nifti image or filename of nifti.
    datatype: type, default: float
        Data type of the output array. For masks, use int.
    nonzeroneighbors : boolean, default=False
        If True, zero datapoints are replaced by largest absolute neihbors voxel data.
        This is useful for masked niftis, where data outside of the mask is zero and you want to pull data
        only from inside the mask.

    Returns
    -------
    np.ndarray of float
        (n,) array with voxel values from nifti file for each triangle center.
    """
    coords_center = np.average(tri_node_coords, axis=1)

    if isinstance(nii, str):
        nii = nib.load(nii)

    # load nifti data
    nii_data = nii.get_fdata().copy()
    affine = nii.affine
    affine_inv = np.linalg.inv(affine)

    # create empty results arrays
    fmri_surf = np.zeros_like(coords_center)

    # pull voxel data for each triangle
    t = tqdm(coords_center)
    t.set_description("Mapping .nii to mesh")
    for i, coord in enumerate(t):
        # we need to apply an array2voxel coordinates transformations
        fmri_surf[i, :] = apply_affine(affine_inv, coord)

    fmri_stats = np.round(fmri_surf).astype(int)
    for dim in range(3):
        fmri_stats[:, dim] = np.clip(fmri_stats[:, dim], 0, nii_data.shape[dim] - 1)
    # print(np.min(fmri_stats[:,0]), np.max(fmri_stats[:,0]))
    # print(np.min(fmri_stats[:, 1]), np.max(fmri_stats[:, 1]))
    # print(np.min(fmri_stats[:,2]), np.max(fmri_stats[:,2]))
    # a = nii_data[fmri_stats[:, 0], fmri_stats[:, 1], fmri_stats[:, 2]].astype(datatype)
    # b = nii_data[fmri_stats[:, 0], fmri_stats[:, 1], fmri_stats[:, 2]].astype(datatype)
    # grab the data

    plusminus = list(itertools.product(*[[1, -1]] * 3))

    mapped_data = nii_data[fmri_stats[:, 0], fmri_stats[:, 1], fmri_stats[:, 2]]
    for zero_idx in np.where(mapped_data == 0)[0]:
        offset_data = []
        zero_arr_coords = fmri_stats[zero_idx]
        # mapped_data_zero = nii_data[zero_arr_coords[0], zero_arr_coords[1], zero_arr_coords[2]][0]
        for x,y,z in plusminus:
            offset_data.append(nii_data[zero_arr_coords[0]+y,
                               zero_arr_coords[1]+y,
                               zero_arr_coords[2]+z][0])
        best_offset = np.argmax(np.abs(offset_data))
        fmri_stats[zero_idx] += plusminus[best_offset]

    return nii_data[fmri_stats[:, 0], fmri_stats[:, 1], fmri_stats[:, 2]].astype(datatype)
