import os
import numpy as np
from shutil import rmtree
from .hclust import clusterTools as CT
from tempfile import mkdtemp
from pathlib import Path
from ..utils import read_bundle, write_bundle, sampling
from .hclust.c_wrappers import fiberDistanceMax, getAffinityGraphFromDistanceMatrix, getAverageLinkHCFromGraphFile
import logging
import time
import joblib
from .ffclust.clustering import split_fibers, parallel_points_clustering, MapClustering,\
    small_clusters_reassignment, get_groups, parallel_group_join_clique
from .ffclust.utils import clusters_fibers, save_clusters, save_clusters_centroids

"""
HClust Sub-Module.
"""

def hierarchical(fiber_input, outfile_dir, MaxDistance_Thr, PartDistance_Thr, variance, clusters_id_dir):
    """
    Run to Hierarchical clustering
    """
    matrix_distance_dir = os.path.join(outfile_dir, 'distance_matrix.bin')
    affinity_graph_dir = os.path.join(outfile_dir,'affinity_graph.txt')
    dendogram_dir = os.path.join(outfile_dir,'dendogram.txt')

    # Distance matrix
    fiberDistanceMax(fiber_input, matrix_distance_dir)

    # affinity graph
    getAffinityGraphFromDistanceMatrix(matrix_distance_dir, affinity_graph_dir, MaxDistance_Thr)

    # Dendogram
    getAverageLinkHCFromGraphFile(affinity_graph_dir, dendogram_dir)

    # Dendogram partition
    wfv=CT.wforest_partition_maxdist_from_graph(dendogram_dir,PartDistance_Thr,True,affinity_graph_dir,variance)
    clusters=wfv.clusters

    ar=open(clusters_id_dir,'wt')
    ar.write(str(clusters))
    ar.close()

    return clusters



def write_fiber_clusters (fiber_input, clusters, final_bundles_dir):

    """Writes the fiber clusters
    """  
    os.makedirs(final_bundles_dir, exist_ok=True)
    
    for i, j in enumerate(clusters):
        write_bundle(os.path.join(final_bundles_dir, f"{i}.bundles"), fiber_input[j])



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 = sorted(p.glob("*.bundles"), key=lambda x: int(x.stem))
    for cluster_path in clusters_paths:
        cluster = read_bundle(str(cluster_path))
        centroids.append(cal_centroide(cluster))
    write_bundle(file_out, centroids)


def hclust(file_in: str, dir_out: str, fiber_thr: int, partition_thr: int, variance: int) -> None:
    """
    Average-link hierarchical agglomerative clustering algorithm which allows finding bundles based on a pairwise fiber distance measure.

    Parameters
    ----------
    file_in : str
        Tractography dataset file in *bundles* format.
    dir_out : str
        Directory to store all the results generated by the algorithm.
    fiber_thr : str
        Maximum distance threshold in *mm*, default: *70*.
    partition_thr : str
        Partition threshold in *mm*, default: *70*.
    variance : str
        Variance squared and provides a similarity scale in *mm*, default: *3600*.

    Returns
    -------
    None

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

    final_bundles : bundles files
        Directory that stores all the fiber clusters identified in separated datasets (*bundles/bundlesdata* files), sampled at 21 points. 
        The file names are labeled with integer numbers ranging from zero to the total number of fiber clusters (N-1) identified.
    final_bundles_allpoints : bundles files    
        Directory that stores all the fiber clusters identified in separated datasets (*bundles/bundlesdata* files). 
        The file names are labeled with integer numbers ranging from zero to the total number of fiber clusters (N-1) identified.
        This directory is only generated when the input tractography dataset file has not been sampled at 21 points. 
    centroids : bundles files
        A directory storing two files corresponding to a tractography dataset in *bundles/bundlesdata* format, 
        containing one centroid for each created cluster. This dataset is named *centroids.bundles/centroids.bundlesdata*, and is sampled at 21 points in the atlas space (MNI).
    bundles_id : text file
        Text file that stores the fiber indexes from the original input tractography dataset file for each detected cluster. 
        Each line in the file correspond to a cluster in correlative order.
    outputs : Several formats
        Temporal directory with intermediate results, such as the distance matrix, the affinity graph, and the dendrogram.

    """
    
    if os.path.exists(dir_out):
        rmtree(dir_out)

    os.mkdir(dir_out)

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

    final_centroids_dir = os.path.join(dir_out, 'centroids')
    final_centroids_file = os.path.join(final_centroids_dir, 'centroids.bundles')
    os.makedirs(final_centroids_dir, exist_ok=True)

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

    clusters_id_dir = os.path.join(dir_out, 'bundles_id.txt')

    num_point_fibers = 1
    data = read_bundle(file_in)
    for fib in data:
        if len(fib) != 21:
            num_point_fibers = 0
            break

    if num_point_fibers == 0:

        fibers21p = os.path.join(outfile_dir, 'fibers_21p.bundles')
        sampling(file_in, fibers21p, 21)
       
        clusters = hierarchical(fibers21p, outfile_dir, fiber_thr, partition_thr, variance, clusters_id_dir)

        final_bundles_allpoints_dir = os.path.join(dir_out, 'final_bundles_allpoints')
        os.makedirs(final_bundles_allpoints_dir, exist_ok=True)

        fibers_allpoints = np.asarray(data, dtype = object)
        write_fiber_clusters (fibers_allpoints, clusters, final_bundles_allpoints_dir)

        fibers_21points = np.asarray(read_bundle(fibers21p))
        write_fiber_clusters (fibers_21points, clusters, final_bundles21p_dir)

        write_centroids(final_bundles21p_dir, final_centroids_file)

    else:

        clusters = hierarchical(file_in, outfile_dir, fiber_thr, partition_thr, variance, clusters_id_dir)   
        
        fibers_input = np.asarray(data)     
        write_fiber_clusters (fibers_input, clusters, final_bundles21p_dir)        
        
        write_centroids(final_bundles21p_dir, final_centroids_file)



def ffclust(file_in: str, dir_out: str, points: [0,3,10,17,20], ks: [200,300,300,300,200], assign_thr: int=6, join_thr: int=6) -> bool:
    """
    Intra-subject clustering algorithm aims to identify compact and homogeneous fiber clusters on a large tractography dataset.

    Parameters
    ----------
    file_in : str
        Tractography data file in *'.bundles'* format.
    dir_out : str
        Directory to store all the results generated by the algorithm.
    points : list()
        Index of the points to be used in the point clustering (K-means), default: 0, 3, 10, 17, 20.
    ks : list()
        Number of clusters to be computed for each point using K-Means, default: 300, 200, 200, 200, 300.
    assign_thr : str
        Maximum distance threshold for the cluster reassignment in *mm*, default: *6*.
    join_thr : str
        Maximum distance threshold for the cluster merge in *mm*, default: *6*.

    Returns
    -------
    None

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

    final_bundles : bundles files
        Directory that stores all the fiber clusters identified in separated datasets (*bundles/bundlesdata* files), sampled at 21 points. 
        The file names are labeled with integer numbers ranging from zero to the total number of fiber clusters (N-1) identified.
    final_bundles_allpoints : bundles files    
        Directory that stores all the fiber clusters identified in separated datasets (*bundles/bundlesdata* files). 
        The file names are labeled with integer numbers ranging from zero to the total number of fiber clusters (N-1) identified.
        This directory is only generated when the input tractography dataset file has not been sampled at 21 points. 
    centroids : bundles files
        A directory storing two files corresponding to a tractography dataset in *bundles/bundlesdata* format, 
        containing one centroid for each created cluster. This dataset is named *centroids.bundles/centroids.bundlesdata*, and is sampled at 21 points in the atlas space (MNI).
    bundles_id : text file
        Text file that stores the fiber indexes from the original input tractography dataset file for each detected cluster. 
        Each line in the file correspond to a cluster in correlative order.
    outputs : Several formats
        Temporal directory with intermediate results, such as the point clusters.       
    """

    if isinstance(assign_thr, str):
        assign_thr = int(assign_thr)
    if isinstance(join_thr, str):
        join_thr = int(join_thr)

    point = np.array(points)
    kss = np.array(ks)

    os.makedirs(dir_out, exist_ok=True)

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

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

    object_dir = os.path.join(dir_out, 'outputs')
    os.makedirs(object_dir, exist_ok=True)
    data=read_bundle(file_in)

    num_point_fibers=1
    for fib in data:
        if len(fib) != 21:
            num_point_fibers=0
            break


    if num_point_fibers == 0:
        fibers21p = os.path.join(object_dir,'fiber_input_21pts.bundles')
        sampling(file_in, fibers21p, 21)
        fibers = np.asarray(read_bundle(fibers21p))
    
    else:
        fibers = np.asarray(data)

    # New logging file each time, changing filemode to a should append instead.
    logging.basicConfig(filename=os.path.join(object_dir,'info.log'), filemode='w', level=logging.INFO)

    t1 = time.time()
    X = split_fibers(fibers[:,point,:],point)
    labels, clusterers = parallel_points_clustering(X=X, ks=kss)
    logging.info('Kmeans Time: {}'.format(time.time() - t1))

    t1 = time.time()
    m = MapClustering()
    map_clusters = m.cluster(fibers, labels)
    logging.info('Map Time: {}'.format(time.time() - t1))

    map_output_filename = os.path.join(object_dir, 'clusters_map.txt')
    save_clusters(dataset=fibers, clusters=map_clusters, filename=map_output_filename)

    with open(os.path.join(object_dir,'stats.txt'), 'w') as f:
        f.write('Number of clusters in map_clusters: {}\n'.format(len(map_clusters)))
        f.write('Number of fibers in map_clusters: {}\n'.format(sum(map_clusters.clusters_sizes())))

    for p,k, clusterer in zip(point, kss, clusterers):
        joblib.dump(clusterer, os.path.join (object_dir, 'clusterer-Point{}-k{}.pkl'.format(p,k)))

  
    t1 = time.time()
    actual_clusters = small_clusters_reassignment(clusters=map_clusters,
                                                  min_size_filter=6,
                                                  max_size_filter=5,
                                                  input_dir = 'segmentation/bundles/result/parallelFastCPU',
                                                  threshold = assign_thr,
                                                  refdata = fibers)

    ident_clusters = dict()
    for i,c in enumerate(actual_clusters):
        ident_clusters[str(c.indices)] = i

    logging.info('Segmentation Time {}'.format(time.time() - t1))

    t1 = time.time()
    point_index = 10
    ngroups = kss[len(kss)//2]

    centroids = np.asarray([x.centroid for x in actual_clusters])
    centroids_points = centroids[:,point_index]
    clusterer=clusterers[len(point)//2]
    labels = clusterer.predict(centroids_points)

    groups = get_groups(labels, ngroups=ngroups)

    final_clusters = parallel_group_join_clique(actual_clusters, groups, fibers, final_bundles21p_dir, dir_out, final_centroids_dir, ident_clusters, object_dir, join_thr)
    final_centroids_filename = os.path.join(object_dir, 'centroids.txt')
    save_clusters_centroids(clusters=final_clusters, filename=final_centroids_filename)

    os.makedirs(final_bundles21p_dir, exist_ok=True)
    clusters_fibers (final_clusters, final_bundles21p_dir, fibers)

    final_bundles_allpoints_dir = os.path.join(dir_out, 'final_bundles_allpoints')
    os.makedirs(final_bundles_allpoints_dir, exist_ok=True)

    if num_point_fibers == 0:
        clusters_fibers (final_clusters, final_bundles_allpoints_dir, np.asarray(data, dtype=object))
    

    return True
