# -*- coding: utf-8 -*-

from copy import deepcopy
import json
import logging
import os
import shutil
from typing import Any, List, Tuple, Type, Union, Optional
import zipfile

import nibabel as nib
from nibabel.affines import voxel_sizes
from nibabel.nifti1 import Nifti1Header, Nifti1Image
from nibabel.orientations import aff2axcodes
from nibabel.streamlines.array_sequence import ArraySequence
from nibabel.streamlines.trk import TrkFile
from nibabel.streamlines.tractogram import Tractogram, LazyTractogram
import numpy as np

from trx.io import get_trx_tmp_dir
from trx.utils import (append_generator_to_dict,
                       close_or_delete_mmap,
                       convert_data_dict_to_tractogram,
                       get_reference_info_wrapper)

try:
    import dipy
    dipy_available = True
except ImportError:
    dipy_available = False


def _append_last_offsets(nib_offsets: np.ndarray, nb_vertices: int) -> np.ndarray:
    """Appends the last element of offsets from header information

    Keyword arguments:
        nib_offsets -- np.ndarray
            Array of offsets with the last element being the start of the last
            streamline (nibabel convention)
        nb_vertices -- int
            Total number of vertices in the streamlines
    Returns:
        Offsets -- np.ndarray (VTK convention)
    """
    def is_sorted(a): return np.all(a[:-1] <= a[1:])
    if not is_sorted(nib_offsets):
        raise ValueError('Offsets must be sorted values.')
    return np.append(nib_offsets, nb_vertices).astype(nib_offsets.dtype)


def _generate_filename_from_data(arr: np.ndarray, filename: str) -> str:
    """Determines the data type from array data and generates the appropriate
    filename

    Keyword arguments:
        arr -- a NumPy array (1-2D, otherwise ValueError raised)
        filename -- the original filename

    Returns:
        An updated filename
    """
    base, ext = os.path.splitext(filename)
    if ext:
        logging.warning("Will overwrite provided extension if needed.")

    dtype = arr.dtype
    dtype = "bit" if dtype == bool else dtype.name

    if arr.ndim == 1:
        new_filename = "{}.{}".format(base, dtype)
    elif arr.ndim == 2:
        dim = arr.shape[-1]
        if dim == 1:
            new_filename = "{}.{}".format(base, dtype)
        else:
            new_filename = "{}.{}.{}".format(base, arr.shape[-1], dtype)
    else:
        raise ValueError("Invalid dimensionality.")

    return new_filename


def _split_ext_with_dimensionality(filename: str) -> Tuple[str, int, str]:
    """Takes a filename and splits it into its components

    Keyword arguments:
        filename -- Input filename

    Returns:
        tuple of strings (basename, dimension, extension)

    """
    basename = os.path.basename(filename)
    split = basename.split(".")

    if len(split) != 2 and len(split) != 3:
        raise ValueError("Invalid filename.")
    basename = split[0]
    ext = ".{}".format(split[-1])
    dim = 1 if len(split) == 2 else split[1]

    _is_dtype_valid(ext)

    return basename, int(dim), ext


def _compute_lengths(offsets: np.ndarray) -> np.ndarray:
    """Compute lengths from offsets

    Keyword arguments:
        offsets -- An np.ndarray of offsets

    Returns:
        lengths -- An np.ndarray of lengths
    """
    if len(offsets) > 0:
        last_elem_pos = _dichotomic_search(offsets)
        lengths = np.ediff1d(offsets)
        if len(lengths) > last_elem_pos:
            lengths[last_elem_pos] = 0
    else:
        lengths = np.array([0])

    return lengths.astype(np.uint32)


def _is_dtype_valid(ext: str) -> bool:
    """Verifies that filename extension is a valid datatype

    Keyword arguments:
        ext -- filename extension

    Returns:
        boolean representing if provided datatype is valid
    """
    if ext.replace(".", "") == "bit":
        return True
    try:
        isinstance(np.dtype(ext.replace(".", "")), np.dtype)
        return True
    except TypeError:
        return False


def _dichotomic_search(
    x: np.ndarray, l_bound: Optional[int] = None, r_bound: Optional[int] = None
) -> int:
    """Find where data of a contiguous array is actually ending

    Keyword arguments:
        x -- np.ndarray of values
        l_bound -- lower bound index for search
        r_bound -- upper bound index for search
    Returns:
        index at which array value is 0 (if possible), otherwise returns -1"""
    if l_bound is None and r_bound is None:
        l_bound = 0
        r_bound = len(x) - 1

    if l_bound == r_bound:
        val = l_bound if x[l_bound] != 0 else -1
        return val

    mid_bound = (l_bound + r_bound + 1) // 2

    if x[mid_bound] == 0:
        return _dichotomic_search(x, l_bound, mid_bound - 1)
    else:
        return _dichotomic_search(x, mid_bound, r_bound)


def _create_memmap(
    filename: str,
    mode: str = "r",
    shape: Tuple = (1,),
    dtype: np.dtype = np.float32,
    offset: int = 0,
    order: str = "C",
) -> np.ndarray:
    """Wrapper to support empty array as memmaps

    Keyword arguments:
        filename -- filename where the empty memmap should be created
        mode -- file open mode (see: np.memmap for options)
        shape -- shape of memmapped np.ndarray
        dtype -- datatype of memmapped np.ndarray
        offset -- offset of the data within the file
        order -- data representation on disk (C or Fortran)

    Returns:
        mmapped np.ndarray or a zero-filled Numpy array if array has a shape of 0
            in the first dimension
    """
    if np.dtype(dtype) == bool:
        filename = filename.replace(".bool", ".bit")

    if shape[0]:
        return np.memmap(
            filename, mode=mode, offset=offset, shape=shape, dtype=dtype,
            order=order
        )
    else:
        if not os.path.isfile(filename):
            f = open(filename, "wb")
            f.close()
        return np.zeros(shape, dtype=dtype)


def load(input_obj: str, check_dpg: bool = True) -> Type["TrxFile"]:
    """Load a TrxFile (compressed or not)

    Keyword arguments:
    input_obj -- A directory name or filepath to the trx data
    check_dpg -- Boolean denoting if group metadata should be checked

    Returns:
        TrxFile object representing the read data
    """
    # TODO Check if 0 streamlines, then 0 vertices is expected (vice-versa)
    # TODO 4x4 affine matrices should contains values (no all-zeros)
    # TODO 3x1 dimensions array should contains values at each position (int)
    if os.path.isfile(input_obj):
        was_compressed = False
        with zipfile.ZipFile(input_obj, "r") as zf:
            for info in zf.infolist():
                if info.compress_type != 0:
                    was_compressed = True
                    break
        if was_compressed:
            with zipfile.ZipFile(input_obj, "r") as zf:
                tmp_dir = get_trx_tmp_dir()
                zf.extractall(tmp_dir.name)
                trx = load_from_directory(tmp_dir.name)
                trx._uncompressed_folder_handle = tmp_dir
                logging.info(
                    "File was compressed, call the close() function before "
                    "exiting."
                )
        else:
            trx = load_from_zip(input_obj)
    elif os.path.isdir(input_obj):
        trx = load_from_directory(input_obj)
    else:
        raise ValueError("File/Folder does not exist")

    # Example of robust check for metadata
    if check_dpg:
        for dpg in trx.data_per_group.keys():
            if dpg not in trx.groups.keys():
                raise ValueError(
                    "An undeclared group ({}) has " "data_per_group.".format(
                        dpg)
                )
    return trx


def load_from_zip(filename: str) -> Type["TrxFile"]:
    """Load a TrxFile from a single zipfile. Note: does not work with
    compressed zipfiles

    Keyword arguments:
    filename -- path of the zipped TrxFile

    Returns:
        TrxFile representing the read data
    """
    with zipfile.ZipFile(filename, mode="r") as zf:
        with zf.open("header.json") as zf_header:
            header = json.load(zf_header)
            header["VOXEL_TO_RASMM"] = np.reshape(
                header["VOXEL_TO_RASMM"], (4, 4)
            ).astype(np.float32)
            header["DIMENSIONS"] = np.array(
                header["DIMENSIONS"], dtype=np.uint16)

        files_pointer_size = {}
        for zip_info in zf.filelist:
            elem_filename = zip_info.filename
            _, ext = os.path.splitext(elem_filename)
            if ext == ".json" or zip_info.is_dir():
                continue

            if not _is_dtype_valid(ext):
                continue
                raise ValueError(
                    "The dtype {} is not supported".format(elem_filename))

            if ext == ".bit":
                ext = ".bool"

            mem_adress = zip_info.header_offset + len(zip_info.FileHeader())
            dtype_size = np.dtype(ext[1:]).itemsize
            size = zip_info.file_size / dtype_size

            if len(zip_info.extra):
                mem_adress -= len(zip_info.extra)

            if size.is_integer():
                files_pointer_size[elem_filename] = mem_adress, int(size)
            else:
                raise ValueError("Wrong size or datatype")

    return TrxFile._create_trx_from_pointer(
        header, files_pointer_size, root_zip=filename
    )


def load_from_directory(directory: str) -> Type["TrxFile"]:
    """Load a TrxFile from a folder containing memmaps

    Keyword arguments:
    filename -- path of the zipped TrxFile

    Returns:
        TrxFile representing the read data
    """

    directory = os.path.abspath(directory)
    with open(os.path.join(directory, "header.json")) as header:
        header = json.load(header)
        header["VOXEL_TO_RASMM"] = np.reshape(header["VOXEL_TO_RASMM"],
                                              (4, 4)).astype(np.float32)
        header["DIMENSIONS"] = np.array(header["DIMENSIONS"], dtype=np.uint16)
    files_pointer_size = {}
    for root, dirs, files in os.walk(directory):
        for name in files:
            elem_filename = os.path.join(root, name)
            _, ext = os.path.splitext(elem_filename)
            if ext == ".json":
                continue

            if not _is_dtype_valid(ext):
                raise ValueError(
                    "The dtype of {} is not supported".format(elem_filename)
                )

            if ext == ".bit":
                ext = ".bool"

            dtype_size = np.dtype(ext[1:]).itemsize
            size = os.path.getsize(elem_filename) / dtype_size
            if size.is_integer():
                files_pointer_size[elem_filename] = 0, int(size)
            elif os.path.getsize(elem_filename) == 1:
                files_pointer_size[elem_filename] = 0, 0
            else:
                raise ValueError("Wrong size or datatype")

    return TrxFile._create_trx_from_pointer(header, files_pointer_size,
                                            root=directory)


def concatenate(
    trx_list: List["TrxFile"],
    delete_dpv: bool = False,
    delete_dps: bool = False,
    delete_groups: bool = False,
    check_space_attributes: bool = True,
    preallocation: bool = False,
) -> "TrxFile":
    """Concatenate multiple TrxFile together, support preallocation

    Keyword arguments:
        trx_list -- A list containing TrxFiles to concatenate
        delete_dpv -- Delete dpv keys that do not exist in all the provided
            TrxFiles
        delete_dps -- Delete dps keys that do not exist in all the provided
            TrxFile
        delete_groups -- Delete all the groups that currently exist in the
            TrxFiles
        check_space_attributes -- Verify that dimensions and size of data are
            similar between all the TrxFiles
        preallocation -- Preallocated TrxFile has already been generated and
            is the first element in trx_list
            (Note: delete_groups must be set to True as well)

    Returns:
        TrxFile representing the concatenated data

    """
    trx_list = [
        curr_trx for curr_trx in trx_list if curr_trx.header["NB_STREAMLINES"] > 0
    ]
    if len(trx_list) == 0:
        logging.warning("Inputs of concatenation were empty.")
        return TrxFile()

    ref_trx = trx_list[0]
    all_dps = []
    all_dpv = []
    for curr_trx in trx_list:
        all_dps.extend(list(curr_trx.data_per_streamline.keys()))
        all_dpv.extend(list(curr_trx.data_per_vertex.keys()))
    all_dps, all_dpv = set(all_dps), set(all_dpv)

    if check_space_attributes:
        for curr_trx in trx_list[1:]:
            if not np.allclose(
                ref_trx.header["VOXEL_TO_RASMM"], curr_trx.header["VOXEL_TO_RASMM"]
            ) or not np.array_equal(
                ref_trx.header["DIMENSIONS"], curr_trx.header["DIMENSIONS"]
            ):
                raise ValueError("Wrong space attributes.")

    if preallocation and not delete_groups:
        raise ValueError(
            "Groups are variables, cannot be handled with " "preallocation"
        )

    # Verifying the validity of fixed-size arrays, coherence between inputs
    for curr_trx in trx_list:
        for key in all_dpv:
            if key not in ref_trx.data_per_vertex.keys() \
                    or key not in curr_trx.data_per_vertex.keys():
                if not delete_dpv:
                    logging.debug(
                        "{} dpv key does not exist in all TrxFile.".format(key)
                    )
                    raise ValueError(
                        "TrxFile must be sharing identical dpv " "keys.")
            elif (
                ref_trx.data_per_vertex[key]._data.dtype
                != curr_trx.data_per_vertex[key]._data.dtype
            ):
                logging.debug(
                    "{} dpv key is not declared with the same dtype "
                    "in all TrxFile.".format(key)
                )
                raise ValueError("Shared dpv key, has different dtype.")

    for curr_trx in trx_list:
        for key in all_dps:
            if key not in ref_trx.data_per_streamline.keys() \
                    or key not in curr_trx.data_per_streamline.keys():
                if not delete_dps:
                    logging.debug(
                        "{} dps key does not exist in all " "TrxFile.".format(
                            key)
                    )
                    raise ValueError(
                        "TrxFile must be sharing identical dps " "keys.")
            elif (
                ref_trx.data_per_streamline[key].dtype
                != curr_trx.data_per_streamline[key].dtype
            ):
                logging.debug(
                    "{} dps key is not declared with the same dtype "
                    "in all TrxFile.".format(key)
                )
                raise ValueError("Shared dps key, has different dtype.")

    all_groups_len = {}
    all_groups_dtype = {}
    # Variable-size arrays do not have to exist in all TrxFile
    if not delete_groups:
        for trx_1 in trx_list:
            for group_key in trx_1.groups.keys():
                # Concatenating groups together
                if group_key in all_groups_len:
                    all_groups_len[group_key] += len(trx_1.groups[group_key])
                else:
                    all_groups_len[group_key] = len(trx_1.groups[group_key])
                if (
                    group_key in all_groups_dtype
                    and trx_1.groups[group_key].dtype != all_groups_dtype[group_key]
                ):
                    raise ValueError("Shared group key, has different dtype.")
                else:
                    all_groups_dtype[group_key] = trx_1.groups[group_key].dtype

    # Once the checks are done, actually concatenate
    to_concat_list = trx_list[1:] if preallocation else trx_list
    if not preallocation:
        nb_vertices = 0
        nb_streamlines = 0
        for curr_trx in to_concat_list:
            curr_strs_len, curr_pts_len = curr_trx._get_real_len()
            nb_streamlines += curr_strs_len
            nb_vertices += curr_pts_len

        new_trx = TrxFile(
            nb_vertices=nb_vertices, nb_streamlines=nb_streamlines,
            init_as=ref_trx
        )
        if delete_dps:
            new_trx.data_per_streamline = {}
        if delete_dpv:
            new_trx.data_per_vertex = {}
        if delete_groups:
            new_trx.groups = {}

        tmp_dir = new_trx._uncompressed_folder_handle.name

        # When memory is allocated on the spot, groups and data_per_group can
        # be concatenated together
        for group_key in all_groups_len.keys():
            if not os.path.isdir(os.path.join(tmp_dir, "groups/")):
                os.mkdir(os.path.join(tmp_dir, "groups/"))
            dtype = all_groups_dtype[group_key]
            group_filename = os.path.join(
                tmp_dir, "groups/" "{}.{}".format(group_key, dtype.name)
            )
            group_len = all_groups_len[group_key]
            new_trx.groups[group_key] = _create_memmap(
                group_filename, mode="w+", shape=(group_len,), dtype=dtype
            )
            if delete_groups:
                continue
            pos = 0
            count = 0
            for curr_trx in trx_list:
                curr_len = len(curr_trx.groups[group_key])
                new_trx.groups[group_key][pos: pos + curr_len] = \
                    curr_trx.groups[group_key] + count
                pos += curr_len
                count += curr_trx.header["NB_STREAMLINES"]

        strs_end, pts_end = 0, 0
    else:
        new_trx = ref_trx
        strs_end, pts_end = new_trx._get_real_len()

    for curr_trx in to_concat_list:
        # Copy the TrxFile fixed-size info (the right chunk)
        strs_end, pts_end = new_trx._copy_fixed_arrays_from(
            curr_trx, strs_start=strs_end, pts_start=pts_end
        )
    return new_trx


def save(
    trx: "TrxFile", filename: str, compression_standard: Any = zipfile.ZIP_STORED
) -> None:
    """Save a TrxFile (compressed or not)

    Keyword arguments:
        trx -- The TrxFile to save
        filename -- The path to save the TrxFile to
        compression_standard -- The compression standard to use, as defined by
            the ZipFile library
    """
    _, ext = os.path.splitext(filename)
    if ext not in [".zip", ".trx", ""]:
        raise ValueError("Unsupported extension.")

    copy_trx = trx.deepcopy()
    copy_trx.resize()
    tmp_dir_name = copy_trx._uncompressed_folder_handle.name
    if ext in [".zip", ".trx"]:
        zip_from_folder(tmp_dir_name, filename, compression_standard)
    else:
        if os.path.isdir(filename):
            shutil.rmtree(filename)
        shutil.copytree(tmp_dir_name, filename)
    copy_trx.close()


def zip_from_folder(
    directory: str, filename: str, compression_standard: Any = zipfile.ZIP_STORED
) -> None:
    """Utils function to zip on-disk memmaps

    Keyword arguments
        directory -- The path to the on-disk memmap
        filename -- The path where the zip file should be created
        compression_standard -- The compression standard to use, as defined by
            the ZipFile library

    """
    with zipfile.ZipFile(filename, mode="w", compression=compression_standard) as zf:
        for root, _, files in os.walk(directory):
            for name in files:
                curr_filename = os.path.join(root, name)
                tmp_filename = curr_filename.replace(directory, "")[1:]
                zf.write(curr_filename, tmp_filename)


class TrxFile:
    """Core class of the TrxFile"""

    header: dict
    streamlines: Type[ArraySequence]
    groups: dict
    data_per_streamline: dict
    data_per_vertex: dict
    data_per_group: dict

    def __init__(
        self,
        nb_vertices: Optional[int] = None,
        nb_streamlines: Optional[int] = None,
        init_as: Optional[Type["TrxFile"]] = None,
        reference: Union[
            str,
            dict,
            Type[Nifti1Image],
            Type[TrkFile],
            Type[Nifti1Header],
            None,
        ] = None,
    ) -> None:
        """Initialize an empty TrxFile, support preallocation

        Keyword Arguments:
            nb_vertices -- The number of vertices to use in the new TrxFile
            nb_streamlines -- The number of streamlines in the new TrxFile
            init_as -- A TrxFile to use as reference
            reference -- A Nifti or Trk file/obj to use as reference
        """
        if init_as is not None:
            affine = init_as.header["VOXEL_TO_RASMM"]
            dimensions = init_as.header["DIMENSIONS"]
        elif reference is not None:
            affine, dimensions, _, _ = get_reference_info_wrapper(reference)
        else:
            logging.debug(
                "No reference provided, using blank space "
                "attributes, please update them later."
            )
            affine = np.eye(4).astype(np.float32)
            dimensions = np.array([1, 1, 1], dtype=np.uint16)

        if nb_vertices is None and nb_streamlines is None:
            if init_as is not None:
                raise ValueError(
                    "Cant use init_as without declaring "
                    "nb_vertices AND nb_streamlines"
                )
            logging.debug("Intializing empty TrxFile.")
            self.header = {}
            # Using the new format default type
            tmp_strs = ArraySequence()
            tmp_strs._data = tmp_strs._data.astype(np.float32)
            tmp_strs._offsets = tmp_strs._offsets.astype(np.uint32)
            tmp_strs._lengths = tmp_strs._lengths.astype(np.uint32)
            self.streamlines = tmp_strs
            self.groups = {}
            self.data_per_streamline = {}
            self.data_per_vertex = {}
            self.data_per_group = {}
            self._uncompressed_folder_handle = None

            nb_vertices = 0
            nb_streamlines = 0

        elif nb_vertices is not None and nb_streamlines is not None:
            logging.debug(
                "Preallocating TrxFile with size {} streamlines"
                "and {} vertices.".format(nb_streamlines, nb_vertices)
            )
            trx = self._initialize_empty_trx(
                nb_streamlines, nb_vertices, init_as=init_as
            )
            self.__dict__ = trx.__dict__
        else:
            raise ValueError(
                "You must declare both nb_vertices AND " "NB_STREAMLINES")

        self.header["VOXEL_TO_RASMM"] = affine
        self.header["DIMENSIONS"] = dimensions
        self.header["NB_VERTICES"] = nb_vertices
        self.header["NB_STREAMLINES"] = nb_streamlines
        self._copy_safe = True

    def __str__(self) -> str:
        """Generate the string for printing"""
        affine = np.array(self.header["VOXEL_TO_RASMM"], dtype=np.float32)
        dimensions = np.array(self.header["DIMENSIONS"], dtype=np.uint16)
        vox_sizes = np.array(voxel_sizes(affine), dtype=np.float32)
        vox_order = "".join(aff2axcodes(affine))

        text = "VOXEL_TO_RASMM: \n{}".format(
            np.array2string(affine, formatter={
                            "float_kind": lambda x: "%.6f" % x})
        )
        text += "\nDIMENSIONS: {}".format(np.array2string(dimensions))
        text += "\nVOX_SIZES: {}".format(
            np.array2string(vox_sizes, formatter={
                            "float_kind": lambda x: "%.2f" % x})
        )
        text += "\nVOX_ORDER: {}".format(vox_order)

        strs_size = self.header["NB_STREAMLINES"]
        pts_size = self.header["NB_VERTICES"]
        strs_len, pts_len = self._get_real_len()

        if strs_size != strs_len or pts_size != pts_len:
            text += "\nstreamline_size: {}".format(strs_size)
            text += "\nvertex_size: {}".format(pts_size)

        text += "\nstreamline_count: {}".format(strs_len)
        text += "\nvertex_count: {}".format(pts_len)
        text += "\ndata_per_vertex keys: {}".format(
            list(self.data_per_vertex.keys()))
        text += "\ndata_per_streamline keys: {}".format(
            list(self.data_per_streamline.keys())
        )

        text += "\ngroups keys: {}".format(list(self.groups.keys()))
        for group_key in self.groups.keys():
            if group_key in self.data_per_group:
                text += "\ndata_per_groups ({}) keys: {}".format(
                    group_key, list(self.data_per_group[group_key].keys())
                )

        text += "\ncopy_safe: {}".format(self._copy_safe)

        return text

    def __len__(self) -> int:
        """Define the length of the object"""
        return len(self.streamlines)

    def __getitem__(self, key) -> Any:
        """Slice all data in a consistent way"""
        if isinstance(key, int):
            if key < 0:
                key += len(self)
            key = [key]
        elif isinstance(key, slice):
            key = [ii for ii in range(*key.indices(len(self)))]

        return self.select(key, keep_group=False)

    def __deepcopy__(self) -> Type["TrxFile"]:
        return self.deepcopy()

    def deepcopy(self) -> Type["TrxFile"]:
        """Create a deepcopy of the TrxFile

        Returns
            A deepcopied TrxFile of the current TrxFile
        """
        tmp_dir = get_trx_tmp_dir()
        out_json = open(os.path.join(tmp_dir.name, "header.json"), "w")
        tmp_header = deepcopy(self.header)

        if not isinstance(tmp_header["VOXEL_TO_RASMM"], list):
            tmp_header["VOXEL_TO_RASMM"] = tmp_header["VOXEL_TO_RASMM"].tolist()
        if not isinstance(tmp_header["DIMENSIONS"], list):
            tmp_header["DIMENSIONS"] = tmp_header["DIMENSIONS"].tolist()

        # tofile() alway write in C-order
        if not self._copy_safe:
            to_dump = self.streamlines.copy()._data
            tmp_header["NB_STREAMLINES"] = len(self.streamlines)
            tmp_header["NB_VERTICES"] = len(to_dump)
        else:
            to_dump = self.streamlines._data
        json.dump(tmp_header, out_json)
        out_json.close()

        positions_filename = _generate_filename_from_data(
            to_dump, os.path.join(tmp_dir.name, "positions")
        )
        to_dump.tofile(positions_filename)

        if not self._copy_safe:
            to_dump = _append_last_offsets(self.streamlines.copy()._offsets,
                                           self.header["NB_VERTICES"])
        else:
            to_dump = _append_last_offsets(self.streamlines._offsets,
                                           self.header["NB_VERTICES"])
        offsets_filename = _generate_filename_from_data(
            self.streamlines._offsets, os.path.join(tmp_dir.name, "offsets")
        )
        to_dump.tofile(offsets_filename)

        if len(self.data_per_vertex.keys()) > 0:
            os.mkdir(os.path.join(tmp_dir.name, "dpv/"))
        for dpv_key in self.data_per_vertex.keys():
            if not self._copy_safe:
                to_dump = self.data_per_vertex[dpv_key].copy()._data
            else:
                to_dump = self.data_per_vertex[dpv_key]._data

            dpv_filename = _generate_filename_from_data(
                to_dump, os.path.join(tmp_dir.name, "dpv/", dpv_key)
            )
            to_dump.tofile(dpv_filename)

        if len(self.data_per_streamline.keys()) > 0:
            os.mkdir(os.path.join(tmp_dir.name, "dps/"))
        for dps_key in self.data_per_streamline.keys():
            to_dump = self.data_per_streamline[dps_key]
            dps_filename = _generate_filename_from_data(
                to_dump, os.path.join(tmp_dir.name, "dps/", dps_key)
            )
            to_dump.tofile(dps_filename)

        if len(self.groups.keys()) > 0:
            os.mkdir(os.path.join(tmp_dir.name, "groups/"))
        for group_key in self.groups.keys():
            to_dump = self.groups[group_key]
            group_filename = _generate_filename_from_data(
                to_dump, os.path.join(tmp_dir.name, "groups/", group_key)
            )
            to_dump.tofile(group_filename)

            if group_key not in self.data_per_group:
                continue
            for dpg_key in self.data_per_group[group_key].keys():
                # Creates 'dpg/' only if required
                if not os.path.isdir(os.path.join(tmp_dir.name, "dpg/")):
                    os.mkdir(os.path.join(tmp_dir.name, "dpg/"))
                if not os.path.isdir(os.path.join(tmp_dir.name, "dpg/", group_key)):
                    os.mkdir(os.path.join(tmp_dir.name, "dpg/", group_key))
                to_dump = self.data_per_group[group_key][dpg_key]
                dpg_filename = _generate_filename_from_data(
                    to_dump, os.path.join(
                        tmp_dir.name, "dpg/", group_key, dpg_key)
                )
                to_dump.tofile(dpg_filename)

        copy_trx = load_from_directory(tmp_dir.name)
        copy_trx._uncompressed_folder_handle = tmp_dir

        return copy_trx

    def _get_real_len(self) -> Tuple[int, int]:
        """Get the real size of data (ignoring zeros of preallocation)

        Returns
            A tuple representing the index of the last streamline and the total
                length of all the streamlines
        """
        if len(self.streamlines._lengths) == 0:
            return 0, 0

        last_elem_pos = _dichotomic_search(self.streamlines._lengths)
        if last_elem_pos != -1:
            strs_end = int(last_elem_pos + 1)
            pts_end = int(np.sum(self.streamlines._lengths[0:strs_end]))
            return strs_end, pts_end

        return 0, 0

    def _copy_fixed_arrays_from(
        self,
        trx: Type["TrxFile"],
        strs_start: int = 0,
        pts_start: int = 0,
        nb_strs_to_copy: Optional[int] = None,
    ) -> Tuple[int, int]:
        """Fill a TrxFile using another and start indexes (preallocation)

        Keyword arguments:
            trx -- TrxFile to copy data from
            strs_start -- The start index of the streamline
            pts_start -- The start index of the point
            nb_strs_to_copy -- The number of streamlines to copy. If not set
                                will copy all

        Returns
            A tuple representing the end of the copied streamlines and end of
                copied points
        """
        if nb_strs_to_copy is None:
            curr_strs_len, curr_pts_len = trx._get_real_len()
        else:
            curr_strs_len = int(nb_strs_to_copy)
            curr_pts_len = np.sum(trx.streamlines._lengths[0:curr_strs_len])
            curr_pts_len = int(curr_pts_len)

        strs_end = strs_start + curr_strs_len
        pts_end = pts_start + curr_pts_len

        if curr_pts_len == 0:
            return strs_start, pts_start

        # Mandatory arrays
        self.streamlines._data[pts_start:pts_end] = \
            trx.streamlines._data[0:curr_pts_len]
        self.streamlines._offsets[strs_start:strs_end] = \
            (trx.streamlines._offsets[0:curr_strs_len] + pts_start)
        self.streamlines._lengths[strs_start:strs_end] = \
            trx.streamlines._lengths[0:curr_strs_len]

        # Optional fixed-sized arrays
        for dpv_key in self.data_per_vertex.keys():
            self.data_per_vertex[dpv_key]._data[
                pts_start:pts_end
            ] = trx.data_per_vertex[dpv_key]._data[0:curr_pts_len]
            self.data_per_vertex[dpv_key]._offsets = self.streamlines._offsets
            self.data_per_vertex[dpv_key]._lengths = self.streamlines._lengths

        for dps_key in self.data_per_streamline.keys():
            self.data_per_streamline[dps_key][
                strs_start:strs_end
            ] = trx.data_per_streamline[dps_key][0:curr_strs_len]

        return strs_end, pts_end

    @staticmethod
    def _initialize_empty_trx(
            nb_streamlines: int, nb_vertices: int,
            init_as: Optional[Type["TrxFile"]] = None) -> Type["TrxFile"]:
        """Create on-disk memmaps of a certain size (preallocation)

        Keyword arguments:
            nb_streamlines -- The number of streamlines that the empty TrxFile
                will be initialized with
            nb_vertices -- The number of vertices that the empty TrxFile will
                be initialized with
            init_as -- A TrxFile to initialize the empty TrxFile with

        Returns:
            An empty TrxFile preallocated with a certain size
        """
        trx = TrxFile()
        tmp_dir = get_trx_tmp_dir()
        logging.info("Temporary folder for memmaps: {}".format(tmp_dir.name))

        trx.header["NB_VERTICES"] = nb_vertices
        trx.header["NB_STREAMLINES"] = nb_streamlines

        if init_as is not None:
            trx.header["VOXEL_TO_RASMM"] = init_as.header["VOXEL_TO_RASMM"]
            trx.header["DIMENSIONS"] = init_as.header["DIMENSIONS"]
            positions_dtype = init_as.streamlines._data.dtype
            offsets_dtype = init_as.streamlines._offsets.dtype
            lengths_dtype = init_as.streamlines._lengths.dtype
        else:
            positions_dtype = np.dtype(np.float16)
            offsets_dtype = np.dtype(np.uint32)
            lengths_dtype = np.dtype(np.uint32)

        logging.debug(
            "Initializing positions with dtype:    {}".format(
                positions_dtype.name)
        )
        logging.debug(
            "Initializing offsets with dtype: {}".format(offsets_dtype.name))
        logging.debug(
            "Initializing lengths with dtype: {}".format(lengths_dtype.name))

        # A TrxFile without init_as only contain the essential arrays
        positions_filename = os.path.join(
            tmp_dir.name, "positions.3.{}".format(positions_dtype.name)
        )
        trx.streamlines._data = _create_memmap(
            positions_filename, mode="w+", shape=(nb_vertices, 3),
            dtype=positions_dtype
        )

        offsets_filename = os.path.join(
            tmp_dir.name, "offsets.{}".format(offsets_dtype.name)
        )
        trx.streamlines._offsets = _create_memmap(
            offsets_filename, mode="w+", shape=(nb_streamlines,),
            dtype=offsets_dtype
        )
        trx.streamlines._lengths = np.zeros(
            shape=(nb_streamlines,), dtype=lengths_dtype
        )

        # Only the structure of fixed-size arrays is copied
        if init_as is not None:
            if len(init_as.data_per_vertex.keys()) > 0:
                os.mkdir(os.path.join(tmp_dir.name, "dpv/"))
            if len(init_as.data_per_streamline.keys()) > 0:
                os.mkdir(os.path.join(tmp_dir.name, "dps/"))

            for dpv_key in init_as.data_per_vertex.keys():
                dtype = init_as.data_per_vertex[dpv_key]._data.dtype
                tmp_as = init_as.data_per_vertex[dpv_key]._data
                if tmp_as.ndim == 1:
                    dpv_filename = os.path.join(
                        tmp_dir.name, "dpv/" "{}.{}".format(
                            dpv_key, dtype.name)
                    )
                    shape = (nb_vertices, 1)
                elif tmp_as.ndim == 2:
                    dim = tmp_as.shape[-1]
                    shape = (nb_vertices, dim)
                    dpv_filename = os.path.join(
                        tmp_dir.name, "dpv/" "{}.{}.{}".format(
                            dpv_key, dim, dtype.name)
                    )
                else:
                    raise ValueError("Invalid dimensionality.")

                logging.debug(
                    "Initializing {} (dpv) with dtype: "
                    "{}".format(dpv_key, dtype.name)
                )
                trx.data_per_vertex[dpv_key] = ArraySequence()
                trx.data_per_vertex[dpv_key]._data = _create_memmap(
                    dpv_filename, mode="w+", shape=shape, dtype=dtype
                )
                trx.data_per_vertex[dpv_key]._offsets = trx.streamlines._offsets
                trx.data_per_vertex[dpv_key]._lengths = trx.streamlines._lengths

            for dps_key in init_as.data_per_streamline.keys():
                dtype = init_as.data_per_streamline[dps_key].dtype
                tmp_as = init_as.data_per_streamline[dps_key]
                if tmp_as.ndim == 1:
                    dps_filename = os.path.join(
                        tmp_dir.name, "dps/" "{}.{}".format(
                            dps_key, dtype.name)
                    )
                    shape = (nb_streamlines,)
                elif tmp_as.ndim == 2:
                    dim = tmp_as.shape[-1]
                    shape = (nb_streamlines, dim)
                    dps_filename = os.path.join(
                        tmp_dir.name, "dps/" "{}.{}.{}".format(
                            dps_key, dim, dtype.name)
                    )
                else:
                    raise ValueError("Invalid dimensionality.")

                logging.debug(
                    "Initializing {} (dps) with and dtype: "
                    "{}".format(dps_key, dtype.name)
                )
                trx.data_per_streamline[dps_key] = _create_memmap(
                    dps_filename, mode="w+", shape=shape, dtype=dtype
                )

        trx._uncompressed_folder_handle = tmp_dir

        return trx

    def _create_trx_from_pointer(
        header: dict,
        dict_pointer_size: dict,
        root_zip: Optional[str] = None,
        root: Optional[str] = None,
    ) -> Type["TrxFile"]:
        """After reading the structure of a zip/folder, create a TrxFile

        Keyword arguments:
            header -- A TrxFile header dictionary which will be used for the
                new TrxFile
            dict_pointer_size -- A dictionary containing the filenames of all
                the files within the TrxFile disk file/folder
            root_zip -- The path of the ZipFile pointer
            root -- The dirname of the ZipFile pointer

        Returns:
            A TrxFile constructer from the pointer provided
        """
        # TODO support empty positions, using optional tag?
        trx = TrxFile()
        trx.header = header
        positions, offsets = None, None
        for elem_filename in dict_pointer_size.keys():
            if root_zip:
                filename = root_zip
            else:
                filename = elem_filename

            folder = os.path.dirname(elem_filename)
            base, dim, ext = _split_ext_with_dimensionality(elem_filename)
            if ext == ".bit":
                ext = ".bool"
            mem_adress, size = dict_pointer_size[elem_filename]

            if root is not None:
                # This is for Unix
                if os.name != 'nt' and folder.startswith(root.rstrip("/")):
                    folder = folder.replace(root, "").lstrip("/")
                # These three are for Windows
                elif os.path.isdir(folder) and os.path.basename(folder) in ['dpv', 'dps', 'groups']:
                    folder = os.path.basename(folder)
                elif os.path.basename(os.path.dirname(folder)) == 'dpg':
                    folder = os.path.join('dpg', os.path.basename(folder))
                else:
                    folder = ''

            # Parse/walk the directory tree
            if base == "positions" and folder == "":
                if size != trx.header["NB_VERTICES"] * 3 or dim != 3:
                    raise ValueError("Wrong data size/dimensionality.")
                positions = _create_memmap(
                    filename,
                    mode="r+",
                    offset=mem_adress,
                    shape=(trx.header["NB_VERTICES"], 3),
                    dtype=ext[1:],
                )
            elif base == "offsets" and folder == "":
                if size != trx.header["NB_STREAMLINES"]+1 or dim != 1:
                    raise ValueError("Wrong offsets size/dimensionality.")
                offsets = _create_memmap(
                    filename,
                    mode="r+",
                    offset=mem_adress,
                    shape=(trx.header["NB_STREAMLINES"]+1,),
                    dtype=ext[1:],
                )
                if offsets[-1] != 0:
                    lengths = _compute_lengths(offsets)
                else:
                    lengths = [0]
            elif folder == "dps":
                nb_scalar = size / trx.header["NB_STREAMLINES"]
                if not nb_scalar.is_integer() or nb_scalar != dim:
                    raise ValueError("Wrong dps size/dimensionality.")
                else:
                    shape = (trx.header["NB_STREAMLINES"], int(nb_scalar))

                trx.data_per_streamline[base] = _create_memmap(
                    filename, mode="r+", offset=mem_adress, shape=shape,
                    dtype=ext[1:]
                )
            elif folder == "dpv":
                nb_scalar = size / trx.header["NB_VERTICES"]
                if not nb_scalar.is_integer() or nb_scalar != dim:
                    raise ValueError("Wrong dpv size/dimensionality.")
                else:
                    shape = (trx.header["NB_VERTICES"], int(nb_scalar))

                trx.data_per_vertex[base] = _create_memmap(
                    filename, mode="r+", offset=mem_adress, shape=shape,
                    dtype=ext[1:]
                )
            elif folder.startswith("dpg"):
                if int(size) != dim:
                    raise ValueError("Wrong dpg size/dimensionality.")
                else:
                    shape = (1, int(size))

                # Handle the two-layers architecture
                data_name = os.path.basename(base)
                sub_folder = os.path.basename(folder)
                if sub_folder not in trx.data_per_group:
                    trx.data_per_group[sub_folder] = {}
                trx.data_per_group[sub_folder][data_name] = _create_memmap(
                    filename, mode="r+", offset=mem_adress, shape=shape,
                    dtype=ext[1:]
                )
            elif folder == "groups":
                # Groups are simply indices, nothing else
                # TODO Crash if not uint?
                if dim != 1:
                    raise ValueError("Wrong group dimensionality.")
                else:
                    shape = (int(size),)
                trx.groups[base] = _create_memmap(
                    filename, mode="r+", offset=mem_adress, shape=shape,
                    dtype=ext[1:]
                )
            else:
                logging.error(
                    "{} is not part of a valid structure.".format(
                        elem_filename)
                )

        # All essential array must be declared
        if positions is not None and offsets is not None:
            trx.streamlines._data = positions
            trx.streamlines._offsets = offsets[:-1]
            trx.streamlines._lengths = lengths
        else:
            raise ValueError("Missing essential data.")

        for dpv_key in trx.data_per_vertex:
            tmp = trx.data_per_vertex[dpv_key]
            trx.data_per_vertex[dpv_key] = ArraySequence()
            trx.data_per_vertex[dpv_key]._data = tmp
            trx.data_per_vertex[dpv_key]._offsets = offsets[:-1]
            trx.data_per_vertex[dpv_key]._lengths = lengths
        return trx

    def resize(
        self,
        nb_streamlines: Optional[int] = None,
        nb_vertices: Optional[int] = None,
        delete_dpg: bool = False,
    ) -> None:
        """Remove the ununsed portion of preallocated memmaps

        Keyword arguments:
            nb_streamlines -- The number of streamlines to keep
            nb_vertices -- The number of vertices to keep
            delete_dpg -- Remove data_per_group when resizing
        """
        if not self._copy_safe:
            raise ValueError("Cannot resize a sliced datasets.")

        strs_end, pts_end = self._get_real_len()

        if nb_streamlines is not None and nb_streamlines < strs_end:
            strs_end = nb_streamlines
            logging.info(
                "Resizing (down) memmaps, less streamlines than it "
                "actually contains."
            )

        if nb_vertices is None:
            pts_end = int(np.sum(self.streamlines._lengths[0:nb_streamlines]))
            nb_vertices = pts_end
        elif nb_vertices < pts_end:
            # Resizing vertices only is too dangerous, not allowed
            logging.warning("Cannot resize (down) vertices for consistency.")
            return

        if nb_streamlines is None:
            nb_streamlines = strs_end

        if (
            nb_streamlines == self.header["NB_STREAMLINES"]
            and nb_vertices == self.header["NB_VERTICES"]
        ):
            logging.debug("TrxFile of the right size, no resizing.")
            return

        trx = self._initialize_empty_trx(
            nb_streamlines, nb_vertices, init_as=self)

        logging.info(
            "Resizing streamlines from size {} to {}".format(
                len(self.streamlines), nb_streamlines
            )
        )
        logging.info(
            "Resizing vertices from size {} to {}".format(
                len(self.streamlines._data), nb_vertices
            )
        )

        # Copy the fixed-sized info from the original TrxFile to the new
        # (resized) one.
        if nb_streamlines < self.header["NB_STREAMLINES"]:
            trx._copy_fixed_arrays_from(self, nb_strs_to_copy=nb_streamlines)
        else:
            trx._copy_fixed_arrays_from(self)

        tmp_dir = trx._uncompressed_folder_handle.name
        if len(self.groups.keys()) > 0:
            os.mkdir(os.path.join(tmp_dir, "groups/"))

        for group_key in self.groups.keys():
            group_dtype = self.groups[group_key].dtype
            group_name = os.path.join(
                tmp_dir, "groups/", "{}.{}".format(group_key, group_dtype.name)
            )
            ori_len = len(self.groups[group_key])

            # Remove groups indices if resizing down
            tmp = self.groups[group_key][self.groups[group_key] < strs_end]
            trx.groups[group_key] = _create_memmap(
                group_name, mode="w+", shape=(len(tmp),), dtype=group_dtype
            )
            logging.debug(
                "{} group went from {} items to {}".format(
                    group_key, ori_len, len(tmp))
            )
            trx.groups[group_key][:] = tmp

        if delete_dpg:
            self.close()
            self.__dict__ = trx.__dict__
            return

        if len(self.data_per_group.keys()) > 0:
            os.mkdir(os.path.join(tmp_dir, "dpg/"))
        for group_key in self.data_per_group:
            if not os.path.isdir(os.path.join(tmp_dir, "dpg/", group_key)):
                os.mkdir(os.path.join(tmp_dir, "dpg/", group_key))
            if group_key not in trx.data_per_group:
                trx.data_per_group[group_key] = {}

            for dpg_key in self.data_per_group[group_key].keys():
                dpg_dtype = self.data_per_group[group_key][dpg_key].dtype
                dpg_filename = _generate_filename_from_data(
                    self.data_per_group[group_key][dpg_key],
                    os.path.join(tmp_dir, "dpg/", group_key, dpg_key),
                )

                shape = self.data_per_group[group_key][dpg_key].shape
                if dpg_key not in trx.data_per_group[group_key]:
                    trx.data_per_group[group_key][dpg_key] = {}
                trx.data_per_group[group_key][dpg_key] = _create_memmap(
                    dpg_filename, mode="w+", shape=shape, dtype=dpg_dtype
                )

                trx.data_per_group[group_key][dpg_key][:] = \
                    self.data_per_group[group_key][dpg_key]

        self.close()
        self.__dict__ = trx.__dict__

    def get_dtype_dict(self):
        """Get the dtype dictionary for the TrxFile

        Returns
            A dictionary containing the dtype for each data element
        """
        dtype_dict = {"positions": self.streamlines._data.dtype,
                      "offsets": self.streamlines._offsets.dtype,
                      "dpv": {}, "dps": {}, "dpg": {}, "groups": {}}

        for key in self.data_per_vertex.keys():
            dtype_dict['dpv'][key] = self.data_per_vertex[key]._data.dtype
        for key in self.data_per_streamline.keys():
            dtype_dict['dps'][key] = self.data_per_streamline[key].dtype

        for group_key in self.data_per_group.keys():
            dtype_dict['groups'][group_key] = self.groups[group_key].dtype

        for group_key in self.data_per_group.keys():
            dtype_dict['dpg'][group_key] = {}
            for dpg_key in self.data_per_group[group_key].keys():
                dtype_dict['dpg'][group_key][dpg_key] = \
                    self.data_per_group[group_key][dpg_key].dtype

        return dtype_dict

    def append(self, obj, extra_buffer: int = 0) -> None:
        curr_dtype_dict = self.get_dtype_dict()
        if dipy_available:
            from dipy.io.stateful_tractogram import StatefulTractogram

        if not isinstance(obj, (TrxFile, Tractogram)) \
                and (dipy_available and not isinstance(obj, StatefulTractogram)):
            raise TypeError(
                "{} is not a supported object type for appending.".format(type(obj)))
        elif isinstance(obj, Tractogram):
            obj = self.from_tractogram(obj, reference=self.header,
                                       dtype_dict=curr_dtype_dict)
        elif dipy_available and isinstance(obj, StatefulTractogram):
            obj = self.from_sft(obj, dtype_dict=curr_dtype_dict)

        self._append_trx(obj, extra_buffer=extra_buffer)

    def _append_trx(self, trx: Type["TrxFile"],
                    extra_buffer: int = 0) -> None:
        """Append a TrxFile to another (support buffer)

        Keyword arguments:
            trx -- The TrxFile to append to the current TrxFile
            extra_buffer -- The additional buffer space required to append data
        """
        strs_end, pts_end = self._get_real_len()

        nb_streamlines = strs_end + trx.header["NB_STREAMLINES"]
        nb_vertices = pts_end + trx.header["NB_VERTICES"]

        if (
            self.header["NB_STREAMLINES"] < nb_streamlines
            or self.header["NB_VERTICES"] < nb_vertices
        ):
            self.resize(
                nb_streamlines=nb_streamlines + extra_buffer,
                nb_vertices=nb_vertices + extra_buffer * 100,
            )
        _ = concatenate([self, trx], preallocation=True, delete_groups=True)

    def get_group(
        self, key: str, keep_group: bool = True, copy_safe: bool = False
    ) -> Type["TrxFile"]:
        """Get a particular group from the TrxFile

        Keyword arguments:
            key -- The group name to select
            keep_group -- Make sure group exists in returned TrxFile
            copy_safe -- Perform a deepcopy

        Returns
            A TrxFile exclusively containing data from said group
        """
        return self.select(self.groups[key], keep_group=keep_group, copy_safe=copy_safe)

    def select(
        self, indices: np.ndarray, keep_group: bool = True, copy_safe: bool = False
    ) -> Type["TrxFile"]:
        """Get a subset of items, always vertices to the same memmaps

        Keyword arguments:
            indices -- The list of indices of elements to return
            keep_group -- Ensure group is returned in output TrxFile
            copy_safe -- Perform a deep-copy

        Returns:
            A TrxFile containing data originating from the selected indices
        """
        indices = np.array(indices, dtype=np.uint32)

        new_trx = TrxFile()
        new_trx._copy_safe = copy_safe
        new_trx.header = deepcopy(self.header)

        if isinstance(indices, np.ndarray) and len(indices) == 0:
            # Even while empty, basic dtype and header must be coherent
            positions_dtype = self.streamlines._data.dtype
            offsets_dtype = self.streamlines._offsets.dtype
            lengths_dtype = self.streamlines._lengths.dtype
            new_trx.streamlines._data = new_trx.streamlines._data.reshape(
                (0, 3)
            ).astype(positions_dtype)
            new_trx.streamlines._offsets = new_trx.streamlines._offsets.astype(
                offsets_dtype
            )
            new_trx.streamlines._lengths = new_trx.streamlines._lengths.astype(
                lengths_dtype
            )
            new_trx.header["NB_VERTICES"] = len(new_trx.streamlines._data)
            new_trx.header["NB_STREAMLINES"] = len(
                new_trx.streamlines._lengths)

            return new_trx.deepcopy() if copy_safe else new_trx

        new_trx.streamlines = (
            self.streamlines[indices].copy(
            ) if copy_safe else self.streamlines[indices]
        )
        for dpv_key in self.data_per_vertex.keys():
            new_trx.data_per_vertex[dpv_key] = (
                self.data_per_vertex[dpv_key][indices].copy()
                if copy_safe
                else self.data_per_vertex[dpv_key][indices]
            )

        for dps_key in self.data_per_streamline.keys():
            new_trx.data_per_streamline[dps_key] = (
                self.data_per_streamline[dps_key][indices].copy()
                if copy_safe
                else self.data_per_streamline[dps_key][indices]
            )

        # Not keeping group is equivalent to the [] operator
        if keep_group:
            logging.warning(
                "Keeping dpg despite affecting the group " "items.")
            for group_key in self.groups.keys():
                # Keep the group indices even when fancy slicing
                index = np.argsort(indices)
                sorted_x = indices[index]
                sorted_index = np.searchsorted(
                    sorted_x, self.groups[group_key])
                yindex = np.take(index, sorted_index, mode="clip")
                mask = indices[yindex] != self.groups[group_key]
                intersect = yindex[~mask]

                if len(intersect) == 0:
                    continue

                new_trx.groups[group_key] = intersect
                if group_key in self.data_per_group:
                    for dpg_key in self.data_per_group[group_key].keys():
                        if group_key not in new_trx.data_per_group:
                            new_trx.data_per_group[group_key] = {}
                        new_trx.data_per_group[group_key][
                            dpg_key
                        ] = self.data_per_group[group_key][dpg_key]

        new_trx.header["NB_VERTICES"] = len(new_trx.streamlines._data)
        new_trx.header["NB_STREAMLINES"] = len(new_trx.streamlines._lengths)
        return new_trx.deepcopy() if copy_safe else new_trx

    @staticmethod
    def from_lazy_tractogram(obj: ["LazyTractogram"], reference,
                             extra_buffer: int = 0,
                             chunk_size: int = 10000,
                             dtype_dict: dict = {'positions': np.float32,
                                                 'offsets': np.uint32,
                                                 'dpv': {}, 'dps': {}}) \
            -> Type["TrxFile"]:
        """Append a TrxFile to another (support buffer)

        Keyword arguments:
            trx -- The TrxFile to append to the current TrxFile
            extra_buffer -- The buffer space between reallocation.
                            This number should be a number of streamlines.
                            Use 0 for no buffer.
            chunk_size -- The number of streamlines to save at a time.
        """

        data = {'strs': [], 'dpv': {}, 'dps': {}}
        concat = None
        count = 0
        iterator = iter(obj)
        while True:
            if count < chunk_size:
                try:
                    i = next(iterator)
                    count += 1
                except StopIteration:
                    obj = convert_data_dict_to_tractogram(data)
                    if concat is None:
                        if len(obj.streamlines) == 0:
                            concat = TrxFile()
                        else:
                            concat = TrxFile.from_tractogram(obj,
                                                             reference=reference,
                                                             dtype_dict=dtype_dict)
                    elif len(obj.streamlines) > 0:
                        curr_obj = TrxFile.from_tractogram(obj,
                                                           reference=reference,
                                                           dtype_dict=dtype_dict)
                        concat.append(curr_obj)
                    break
                append_generator_to_dict(i, data)
            else:
                obj = convert_data_dict_to_tractogram(data)
                if concat is None:
                    concat = TrxFile.from_tractogram(obj,
                                                     reference=reference,
                                                     dtype_dict=dtype_dict)
                else:
                    curr_obj = TrxFile.from_tractogram(obj,
                                                       reference=reference,
                                                       dtype_dict=dtype_dict)
                    concat.append(curr_obj, extra_buffer=extra_buffer)
                data = {'strs': [], 'dpv': {}, 'dps': {}}
                count = 0

        concat.resize()
        return concat

    @staticmethod
    def from_sft(sft, dtype_dict={}):
        """Generate a valid TrxFile from a StatefulTractogram"""

        if len(sft.dtype_dict) > 0:
            dtype_dict = sft.dtype_dict
        if 'dpp' in dtype_dict:
            dtype_dict['dpv'] = dtype_dict.pop('dpp')
        elif len(dtype_dict) == 0:
            dtype_dict = {'positions': np.float32, 'offsets': np.uint32,
                          'dpv': {}, 'dps': {}}

        positions_dtype = dtype_dict['positions']
        offsets_dtype = dtype_dict['offsets']

        if not np.issubdtype(positions_dtype, np.floating):
            logging.warning(
                "Casting positions as {}, considering using a floating point "
                "dtype.".format(positions_dtype))

        if not np.issubdtype(offsets_dtype, np.integer):
            logging.warning(
                "Casting offsets as {}, considering using a integer "
                "dtype.".format(offsets_dtype))

        trx = TrxFile(nb_vertices=len(sft.streamlines._data),
                      nb_streamlines=len(sft.streamlines))
        trx.header = {
            "DIMENSIONS": sft.dimensions.tolist(),
            "VOXEL_TO_RASMM": sft.affine.tolist(),
            "NB_VERTICES": len(sft.streamlines._data),
            "NB_STREAMLINES": len(sft.streamlines),
        }

        old_space = deepcopy(sft.space)
        old_origin = deepcopy(sft.origin)

        # TrxFile are written on disk in RASMM/center convention
        sft.to_rasmm()
        sft.to_center()

        tmp_streamlines = deepcopy(sft.streamlines)

        # Cast the int64 of Nibabel to uint32
        tmp_streamlines._offsets = tmp_streamlines._offsets.astype(
            offsets_dtype)
        tmp_streamlines._data = tmp_streamlines._data.astype(positions_dtype)

        trx.streamlines = tmp_streamlines
        for key in sft.data_per_point:
            dtype_to_use = dtype_dict['dpv'][key] if key in dtype_dict['dpv'] \
                else np.float32
            trx.data_per_vertex[key] = \
                sft.data_per_point[key]
            trx.data_per_vertex[key]._data = \
                sft.data_per_point[key]._data.astype(dtype_to_use)

        for key in sft.data_per_streamline:
            dtype_to_use = dtype_dict['dps'][key] if key in dtype_dict['dps'] \
                else np.float32
            trx.data_per_streamline[key] = sft.data_per_streamline[key].astype(
                dtype_to_use)

        # For safety and for RAM, convert the whole object to memmaps
        tmp_dir = get_trx_tmp_dir()
        save(trx, tmp_dir.name)
        trx.close()
        trx = load_from_directory(tmp_dir.name)
        trx._uncompressed_folder_handle = tmp_dir

        sft.to_space(old_space)
        sft.to_origin(old_origin)
        del tmp_streamlines

        return trx

    @staticmethod
    def from_tractogram(tractogram, reference,
                        dtype_dict={'positions': np.float32,
                                    'offsets': np.uint32,
                                    'dpv': {}, 'dps': {}}):
        """Generate a valid TrxFile from a Nibabel Tractogram"""

        positions_dtype = dtype_dict['positions'] if 'positions' in dtype_dict \
            else np.float32
        offsets_dtype = dtype_dict['offsets'] if 'offsets' in dtype_dict \
            else np.uint32

        if not np.issubdtype(positions_dtype, np.floating):
            logging.warning(
                "Casting positions as {}, considering using a floating point "
                "dtype.".format(positions_dtype))

        if not np.issubdtype(offsets_dtype, np.integer):
            logging.warning(
                "Casting offsets as {}, considering using a integer "
                "dtype.".format(offsets_dtype))

        trx = TrxFile(
            nb_vertices=len(tractogram.streamlines._data),
            nb_streamlines=len(tractogram.streamlines),
        )

        affine, dimensions, _, _ = get_reference_info_wrapper(reference)
        trx.header = {
            "DIMENSIONS": dimensions,
            "VOXEL_TO_RASMM": affine,
            "NB_VERTICES": len(tractogram.streamlines._data),
            "NB_STREAMLINES": len(tractogram.streamlines),
        }

        tmp_streamlines = deepcopy(tractogram.streamlines)

        # Cast the int64 of Nibabel to uint32
        tmp_streamlines._offsets = tmp_streamlines._offsets.astype(
            offsets_dtype)
        tmp_streamlines._data = tmp_streamlines._data.astype(positions_dtype)

        trx.streamlines = tmp_streamlines
        for key in tractogram.data_per_point:
            dtype_to_use = dtype_dict['dpv'][key] if key in dtype_dict['dpv'] \
                else np.float32
            trx.data_per_vertex[key] = \
                tractogram.data_per_point[key]
            trx.data_per_vertex[key]._data = \
                tractogram.data_per_point[key]._data.astype(dtype_to_use)

        for key in tractogram.data_per_streamline:
            dtype_to_use = dtype_dict['dps'][key] if key in dtype_dict['dps'] \
                else np.float32
            trx.data_per_streamline[key] = \
                tractogram.data_per_streamline[key].astype(dtype_to_use)

        # For safety and for RAM, convert the whole object to memmaps
        tmp_dir = get_trx_tmp_dir()
        save(trx, tmp_dir.name)
        trx.close()

        trx = load_from_directory(tmp_dir.name)
        del tmp_streamlines

        return trx

    def to_tractogram(self, resize=False):
        """Convert a TrxFile to a nibabel Tractogram (in RAM)"""
        if resize:
            self.resize()

        trx_obj = self.to_memory()
        tractogram = nib.streamlines.Tractogram([], affine_to_rasmm=np.eye(4))
        tractogram._set_streamlines(trx_obj.streamlines)
        tractogram._data_per_point = trx_obj.data_per_vertex
        tractogram._data_per_streamline = trx_obj.data_per_streamline

        return tractogram

    def to_memory(self, resize: bool = False) -> Type["TrxFile"]:
        """Convert a TrxFile to a RAM representation

        Keyword arguments:
            resize -- Resize TrxFile when converting to RAM representation

        Returns:
            A non memory mapped TrxFile
        """
        if resize:
            self.resize()

        trx_obj = TrxFile()
        trx_obj.header = deepcopy(self.header)
        trx_obj.streamlines = deepcopy(self.streamlines)

        for key in self.data_per_vertex:
            trx_obj.data_per_vertex[key] = deepcopy(self.data_per_vertex[key])

        for key in self.data_per_streamline:
            trx_obj.data_per_streamline[key] = deepcopy(
                self.data_per_streamline[key])

        for key in self.groups:
            trx_obj.groups[key] = deepcopy(self.groups[key])

        for key in self.data_per_group:
            trx_obj.data_per_group[key] = deepcopy(self.data_per_group[key])

        return trx_obj

    def to_sft(self, resize=False):
        """Convert a TrxFile to a valid StatefulTractogram (in RAM)"""
        try:
            from dipy.io.stateful_tractogram import StatefulTractogram, Space
        except ImportError:
            logging.error('Dipy library is missing, cannot convert to '
                          'StatefulTractogram.')
            return None

        affine = np.array(self.header["VOXEL_TO_RASMM"], dtype=np.float32)
        dimensions = np.array(self.header["DIMENSIONS"], dtype=np.uint16)
        vox_sizes = np.array(voxel_sizes(affine), dtype=np.float32)
        vox_order = "".join(aff2axcodes(affine))
        space_attributes = (affine, dimensions, vox_sizes, vox_order)

        if resize:
            self.resize()
        sft = StatefulTractogram(
            deepcopy(self.streamlines),
            space_attributes,
            Space.RASMM,
            data_per_point=deepcopy(self.data_per_vertex),
            data_per_streamline=deepcopy(self.data_per_streamline),
        )
        tmp_dict = self.get_dtype_dict()
        if 'dpv' in tmp_dict:
            tmp_dict['dpp'] = tmp_dict.pop('dpv')
        sft.dtype_dict = self.get_dtype_dict()

        return sft

    def close(self) -> None:
        """Cleanup on-disk temporary folder and initialize an empty TrxFile"""
        if self._uncompressed_folder_handle is not None:
            close_or_delete_mmap(self.streamlines)

            # # Close or delete attributes in dictionaries
            for key in self.data_per_vertex:
                close_or_delete_mmap(self.data_per_vertex[key])

            for key in self.data_per_streamline:
                close_or_delete_mmap(self.data_per_streamline[key])

            for key in self.groups:
                close_or_delete_mmap(self.groups[key])

            for key in self.data_per_group:
                for dpg in self.data_per_group[key]:
                    close_or_delete_mmap(self.data_per_group[key][dpg])

            try:
                self._uncompressed_folder_handle.cleanup()
            except PermissionError:
                logging.error("Windows PermissionError, temporary directory {}"
                              "was not deleted!".format(self._uncompressed_folder_handle.name))
        self.__init__()
        logging.debug("Deleted memmaps and intialized empty TrxFile.")