"""Segmentation Module
"""

import os
import sys
import numpy as np
from pathlib import Path
import shutil
from .fiberseg.c_wrappers import segment
from ..utils import sampling, read_bundle, write_bundle

def is_reversed(cluster):
    base = np.stack([cluster, cluster[:, ::-1]])
    x = base - cluster[None, 0]
    x = np.square(x)
    x = np.sum(x, axis=-1)
    x = np.max(x, axis=-1)
    return x[0] > x[1]

def cal_centroide(cluster):
    cluster = np.asarray(cluster)
    r = is_reversed(cluster)
    if np.any(r):
        c_copy = np.empty_like(cluster)
        c_copy[~r] = cluster[~r]
        c_copy[r] = cluster[r, ::-1]
        cluster = c_copy
    return np.sum(cluster, axis=0) / len(cluster)

def write_centroids(clusters_dir, file_out):
    p = Path(clusters_dir)
    centroids = []
    clusters_paths = p.glob("*.bundles")
    for cluster_path in clusters_paths:
        cluster = read_bundle(str(cluster_path)) 
        c = cal_centroide(cluster)
        centroids.append(c.astype(np.float32))
        
    write_bundle(file_out, centroids)


def fiberseg(file_in: str, subj_name: str, atlas_dir: str, atlas_info: str, dir_out: str) -> None:
    """
    White matter fiber bundle segmentation algorithm based on a multi-subject atlas.

    Parameters
    ----------
    file_in : str
        Tractography dataset file in *bundles* format.
    subj_name : str
        Subject name, used to label the results.
    atlas_dir : str
        Bundle atlas directory, with bundles in separated files, sampled at 21 equidistant points. The bundle atlases provided are in same folders.
    atlas_info : str
        Text file associated to the used atlas, that stores information needed to apply the segmentation algorithm, i.e., a list of the atlas fascicles, containing the name, the segmentation
        threshold (in *mm*) and the size of each fascicle. Note that the segmentation threshold can be adjusted depending on the database to be used.
    dir_out : str
        Directory name to store all the results generated by the algorithm.

    Return
    ------
    None

    Notes
    -----
    This function generates the following files in the specified directory:

    final_bundles : bundles files
        The directory contains individual files for each segmented fascicle in the *bundles/bundlesdata* format, sampled at 21 points or all point, and in the atlas space (MNI). 
        Each file is named using the subject's name followed by the atlas label. 
    centroids : bundles files
        A directory storing two files corresponding to a tractography dataset in *bundles/bundlesdata* format, containing one centroid for each segmented fascicle. 
        This dataset is named *centroids.bundles/centroids.bundlesdata*, and is sampled at 21 points in the atlas space (MNI).
    bundles_id : text files
        The directoty that has a text file for each segmented bundle. A text file of a segmented bundle contains the indices of the fibers in the original  subject's tractography dataset file. 
        The name of each file corresponds to the name of the segmented bundle file with ``'.txt'`` extension.
        These files are used to obtain the segmented fascicles in the subject's original space and with all points of brain fibers.

    """

    id_seg_result= os.path.join(dir_out,'bundles_id')
    os.makedirs(id_seg_result, exist_ok=True)

    has21points = True
    data=read_bundle(file_in)

    for fib in data:
        if len(fib) != 21:
            has21points = False
            break

    final_bundles21p_dir = os.path.join(dir_out, 'final_bundles')
    os.makedirs(final_bundles21p_dir, exist_ok=True)

    centroid_dir= os.path.join(dir_out,'centroids')
    os.makedirs(centroid_dir, exist_ok=True)    

    final_bundles_dir = os.path.join(dir_out, 'final_bundles')
    os.makedirs(final_bundles_dir, exist_ok=True)    

    if not has21points:

        outfile_dir= os.path.join(dir_out, 'outputs')
        os.makedirs(outfile_dir, exist_ok=True)

        fibers21p = os.path.join(outfile_dir,'fibers_21p.bundles')

        sampling(file_in, fibers21p, 21)

        segment(21, fibers21p, subj_name, atlas_dir, atlas_info,
                final_bundles21p_dir, id_seg_result)

        shutil.rmtree(outfile_dir)
               
    else:
        segment(21, file_in, subj_name, atlas_dir, atlas_info,
                final_bundles21p_dir, id_seg_result)  


    write_centroids(os.path.join(dir_out, 'final_bundles'), os.path.join(centroid_dir,'centroids.bundles'))