# -- coding: utf-8 --
"""
HClust Sub-Module.
"""

import gc
import os
import time
import numpy as np
from shutil import rmtree
from . import clusterTools as CT
from tempfile import mkdtemp
from pathlib import Path
from ...utils import read_bundle, write_bundle
from .c_wrappers import fiberDistanceMax, getAffinityGraphFromDistanceMatrix, getAverageLinkHCFromGraphFile
from ...utils import sampling


def hierarchical(fiber_input, outfile_dir, MaxDistance_Threshold):
    """
    Run to Hierarchical clustering
    """
    MatrixDist_output = os.path.join(outfile_dir, 'matrixd.bin')
    affinities_graph_output = os.path.join(outfile_dir,'affin.txt')
    dendogram_output = os.path.join(outfile_dir,'dendogram.txt')

    # Step1. Distance Matrix
    t0 = time.time()
    fiberDistanceMax(fiber_input, MatrixDist_output)
    gc.collect()
    print("Distance Matrix Delay: ", time.time()-t0, "[s]")

    # Step2. Affinities Graph
    t0= time.time()
    getAffinityGraphFromDistanceMatrix(
        MatrixDist_output, affinities_graph_output, MaxDistance_Threshold)
    gc.collect()
    print("Affinities Graph Delay: ", time.time()-t0, "[s]")

    # Step3. Dendogram
    t0= time.time()
    getAverageLinkHCFromGraphFile(
        affinities_graph_output, dendogram_output)
    gc.collect()
    print("Dendogram Delay: ", time.time()-t0, "[s]")


def particional_hierarchical(fiber_input, outfile_dir, PartDistance_Threshold, var, final_bundles_dir):

    """Writes the cluster bundles (bundles format) and fiber indexes per cluster (file '.txt')
    """

    arbfile = os.path.join(outfile_dir, 'dendogram.txt')
    afffile = os.path.join(outfile_dir, 'affin.txt')
    partfile = os.path.join(outfile_dir, 'index_fiberscluster.txt')

    t0= time.time()
    wfv=CT.wforest_partition_maxdist_from_graph(arbfile,PartDistance_Threshold,True,afffile,var)

    clusteres=wfv.clusters

    ar=open(partfile,'wt')
    ar.write(str(clusteres))
    ar.close()

    tractography = np.array(read_bundle(fiber_input))

    os.makedirs(final_bundles_dir, exist_ok=True)

    for i, j in enumerate(clusteres):

        write_bundle(os.path.join(final_bundles_dir, f"{i}.bundles"), tractography[j])

    print("Particional Hierarchical Delay: ", time.time()-t0, "[s]")


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 data 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: *30*.
    partition_thr : str
        Partition threshold in *mm*, default: *40*.
    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:

    Clusters : bundles files
        Directory that stores all the fiber clusters found in different '.bundles' files.
        The file names are labeled with integer numbers ranging from zero to the total number of fiber clusters found.
    Centroids : bundles file
        Directory that contains the centroid for each created cluster in same *'.bundles'* files.
        The firt fiber of bundle centroid is corresponding with centroid calculated for cluster one and so on
    Index of fibers per clusters : text file
        Text file that stores the fiber index input for each of the detected clusters.
        The fiber indexes are extracted from the tractography input. The first line corresponds to cluster zero and so on.

    Examples
    --------
    To test `hclust()`,  download the data from the `link hclust  <https://www.dropbox.com/sh/dt196k65v03eh9m/AABKRW7ad2ssB0N_dpjqK4Dha?dl=1>`_.
    Then, open a Python terminal and run the following commands.

    >>> from phybers.clustering import hclust
    >>> hclust (file_in = 'fibers_test.bundles', dir_out = 'hclust_result', fiber_thr = 30, partition_thr = 40, variance = 3600)

    You will locate the clustering results in the 'hclust_result' directory.
    """
    if os.path.exists(dir_out):
        rmtree(dir_out)

    os.mkdir(dir_out)

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

    final_centroids_dir = os.path.join(dir_out, 'FinalCentroids')
    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, 'outfile')
    os.makedirs(outfile_dir, exist_ok=True)

    np = 0
    data = read_bundle(file_in)
    for i in range(len(data)-1):
        if len(data[i]) != len(data[i+1]):
            np = 21
            break

    if np == 21:

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

        fibers21p = os.path.join(outfile_dir, 'fiberorig_21p.bundles')
        sampling(file_in, fibers21p, np)

        hierarchical(fibers21p, outfile_dir,
                                fiber_thr)
        particional_hierarchical(
            fibers21p, outfile_dir, partition_thr, variance, final_bundles21p_dir)
        write_centroids(final_bundles_dir, final_centroids_file)

    else:
        hierarchical(file_in, outfile_dir, fiber_thr)
        particional_hierarchical(
            file_in, outfile_dir, partition_thr, variance, final_bundles_dir)
        write_centroids(final_bundles_dir, final_centroids_file)
