import argparse
import logging
import time
import os

import numpy as np
import dipy.segment.clustering

import joblib 
import matplotlib.image as mpimg
import phybers.clustering.ffclust.bundleTools as bundleTools
import phybers.clustering.ffclust.clustering as clustering
import phybers.clustering.ffclust.metric as metric
import shutil
from phybers.clustering.ffclust.utils import save_clusters, save_clusters_fibers, save_clusters_centroids

def fastfiber(infile, output_directory, thr_assign=6, thr_join=6):
    points = [0,3,10,17,20]
    points = np.array(points)
    ks = [200,300,300,300,200]
    ks = np.array(ks)
    fibers = np.asarray(bundleTools.read_bundle(infile))
    thr_assign = float(thr_assign)
    thr_join = float(thr_assign)
    
#Root work directory creation, output should be place here.

    if output_directory:
        work_dir = output_directory
    else:
        work_dir = 'output_points{}_ks{}_'.format(str(points), str(ks)) + infile

    if os.path.exists(work_dir):   
      shutil.rmtree(work_dir)

    os.mkdir(work_dir)   

    final_bundles_dir = work_dir + '/FinalBundles'
    os.makedirs(final_bundles_dir, exist_ok=True)

    final_centroids_dir = work_dir + '/FinalCentroids'
    os.makedirs(final_centroids_dir, exist_ok=True)

    map_bundles_dir = work_dir + '/MapBundles'
    os.makedirs(map_bundles_dir, exist_ok=True)

    individual_map_bundles_dir =  map_bundles_dir + '/IndividualBundles'
    os.makedirs(individual_map_bundles_dir, exist_ok=True)

    reassigned_bundles_dir = work_dir + '/ReassignedBundles'
    os.makedirs(reassigned_bundles_dir, exist_ok=True)


    object_dir = work_dir + '/output'
    os.makedirs(object_dir, exist_ok=True)

    map_output_filename = object_dir + '/clusters_map.txt'

    #New logging file each time, changing filemode to a should append instead
    logging.basicConfig(filename=work_dir+'/info.log', filemode='w', level=logging.INFO)
    
    #if __name__ == 'mainFFClust.py':
    if 1 == 1:
        print('Dentro de if')
        pathname = os.path.dirname(__file__)
        colormap = mpimg.imread(pathname + '/colors256.jpg')[0]
        t1 = time.time()
        X = clustering.split_fibers(fibers[:,points,:],points)
        labels, clusterers = clustering.parallel_points_clustering(X=X, ks=ks)
        logging.info('Tiempo Kmeans: {}'.format(time.time() - t1))

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

        final_output_filename = object_dir + '/map_clusters.txt' 
        save_clusters(dataset=fibers, clusters=map_clusters, filename=map_output_filename)

        with open(work_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(points, ks, clusterers):
            joblib.dump(clusterer, object_dir + '/clusterer-Point{}-k{}.pkl'.format(p,k))

        """Write .bundles .bundlesdata of small (< 1 fiber) and long centroids (> 1 fiber) of clusters
         and reassign small clasters to large clusters.
         Input:
          ouput: path of output file directory to write the clusters
          min_size_filter = minimum size to get the largest clusters
          max_size_filter = maximum size to get the clusters with only 1 fiber
          input_dir = path of the input dir that contains the results of segmetation/reassignation
          threshold = minimum distance for the segmentation's method
         """
        #Compenzando la reasignación
        t1 = time.time()
        actual_clusters = clustering.small_clusters_reassignment(clusters=map_clusters,
                                                              min_size_filter=6,
                                                              max_size_filter=5,
                                                              input_dir = 'segmentation/bundles/result/parallelFastCPU',
                                                              output_dir=individual_map_bundles_dir,
                                                              reassignment_dir=reassigned_bundles_dir,
                                                              threshold = thr_assign,
                                                              refdata = fibers)
	    
        #actual_clusters = map_clusters.get_large_clusters(6)
     
        ident_clusters={}
        for i,c in enumerate(actual_clusters):
            ident_clusters[str(c.indices)] = i

        logging.info('Tiempo Segmentacion {}'.format(time.time() - t1))

        t1 = time.time()
        point_index = 10
        ngroups = ks[len(ks)//2]
        
        centroids = np.asarray([x.centroid for x in actual_clusters])
        centroids_points = centroids[:,point_index]
        clusterer=clusterers[len(points)//2]
        labels = clusterer.predict(centroids_points)
        
        groups = clustering.get_groups(labels, ngroups=ngroups)
        joined_clusters = clustering.parallel_group_join_clique(actual_clusters, groups, fibers,final_bundles_dir,final_centroids_dir,ident_clusters,object_dir,thr_join)
        

        final_clusters = joined_clusters
        final_centroids_filename = object_dir + '/final_centroids.txt' 
        save_clusters_centroids(clusters=final_clusters, filename=final_centroids_filename)


