"""

Cluster correlations from correlations.py based on hierarchical linkage.

correlations.py outputs a tuple of

[object 1, object 2, pearson correlation, p value]

Here, we use 1-pearson correlation as our distance

"""

import os
import sys
import math
import argparse

import numpy as np
import scipy.cluster.hierarchy as sch


def parse_text_file(tf, pearsoncol=2, insep="\t"):
    """
    Parse a text file and return an n-choose-2 array of the elements. The array returned has the distance from the first
    element to all other elements, and then the second element to n-1 elements (all but the first), and then the
    third element to n-2 elements (all but the first & second) and so on.
    :param tf: Text file with [a, b, pearson correlation, p-value] (e.g. the output from correlations.py)
    :type tf: str
    :param pearsoncol: the zero indexed column of the pearson correlation score
    :type numcols: int
    :param insep: Input separator. Default = tab
    :type insep: str
    :return: n-choose-2 array of the data and the list of the keys in order
    :rtype: np.array, list
    """

    data = {}
    ks = set()
    with open(tf, 'r') as fin:
        for l in fin:
            p=l.strip().split(insep)
            if len(p) <= pearsoncol:
                sys.stderr.write(f"ERROR: {l} does not have enough entries for {pearsoncol} to be the score. Skipped\n")
                continue
            ks.add(p[0])
            ks.add(p[1])
            if p[0] not in data:
                data[p[0]]={}
            if p[1] not in data:
                data[p[1]] = {}
            val = float(p[pearsoncol])
            if val > 1:
                val = 1
            data[p[0]][p[1]] = val
            data[p[1]][p[0]] = val

    allkeys = list(ks)
    allkeys.sort()
    # our new array will have (len(allkeys)) choose 2 elements
    nal = math.comb(len(allkeys), 2)
    nct = np.empty(nal)
    p = 0
    for i in range(len(allkeys)):
        for j in range(i + 1, len(allkeys)):
            if allkeys[j] not in  data[allkeys[i]]:
                data[allkeys[i]][allkeys[j]] = 0
            nct[p] = 1 - data[allkeys[i]][allkeys[j]]
            p += 1
    return nct, allkeys

def generate_clusters(matrix, idlist, jsonout, noclust=100, print_singles=False):

    L = sch.linkage(matrix, method='average')

    print("Threshold\tNumber of clusters\tSize of largest cluster")
    with open(jsonout, 'w') as out:
        out.write("[\n")
        for ele in range(noclust + 1):
            threshold = ele / noclust
            ind = sch.fcluster(L, threshold, 'distance')
            uniqs, counts = np.unique(ind, return_counts=True)
            freqs = {}
            for idx, u in enumerate(uniqs):
                freqs[u] = counts[idx]

            clusters = {}
            for idx, j in enumerate(ind):
                if freqs[j] == 1 and not print_singles:
                    continue
                if j not in clusters:
                    clusters[j] = []
                clusters[j].append(idlist[idx])

            # res = [i idlist[x] for x in ind]
            # res = [x for x in ind]
            # out.write(f"{{cluster_id : {i}, largest_cluster : {max(counts)}, num_clusters: {uniqs.shape[0]}, clusters: {res}}},\n")
            singles = np.isin(counts, [1]).sum()
            out.write(
                f"{{cluster_id : {ele}, threshold: {threshold}, largest_cluster : {max(counts)}, num_clusters: {uniqs.shape[0]}, num_singleton_clusters: {singles}, clusters: {clusters}}},\n")
            print(f"{threshold}\t{uniqs.shape[0]}\t{max(counts)}")
        out.write("]\n")


def generate_a_cluster(matrix, idlist, jsonout, threshold=0.05, print_singles=False):
    """
    Generate a single cluster given a threshold
    """

    L = sch.linkage(matrix, method='average')
    with open(jsonout, 'w') as out:
        ind = sch.fcluster(L, threshold, 'distance')
        uniqs, counts = np.unique(ind, return_counts=True)
        freqs = {}
        for idx, u in enumerate(uniqs):
            freqs[u] = counts[idx]

        clusters = {}
        for idx, j in enumerate(ind):
            if freqs[j] == 1 and not print_singles:
                continue
            if j not in clusters:
                clusters[j] = []
            clusters[j].append(idlist[idx])

        singles = np.isin(counts, [1]).sum()
        out.write(
            f"{{threshold: {threshold}, largest_cluster : {max(counts)}, num_clusters: {uniqs.shape[0]}, num_singleton_clusters: {singles}, clusters: {clusters}}},\n")
        print(f"{threshold}\t{uniqs.shape[0]}\t{max(counts)}")
        out.write("\n")



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Cluster genes based on Pearson correlation")
    parser.add_argument('-f', '--file', help='file with [a, b, distance]', required=True)
    parser.add_argument('-o', '--output', help='clusters output file name. We print them out in json format', required=True)
    parser.add_argument('-p', '--pearsoncol', help='0 indexed column in input file with the pearson score. Default = 2', type=int, default=2)
    parser.add_argument('-s', '--separator', help='Input separator. Default = tab', default="\t", type=str)
    parser.add_argument('--singles', help='print clusters with one element in them', action='store_true')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-t', '--threshold', help='clustering threshold to print a single cluster', type=float)
    group.add_argument('-n', '--numclust', help='number of clusters to print with a range of thresholds (default=100)', type=int, default=100)

    args = parser.parse_args()

    matrix, idlist = parse_text_file(args.file, args.pearsoncol, args.separator)
    if args.threshold:
        generate_a_cluster(matrix, idlist, args.output, args.threshold, args.singles)
    elif args.numclust:
        generate_clusters(matrix, idlist, args.output, args.numclust, args.singles)

@@ -3528,32 +3528,43 @@
         print(f%22
+Threshold: 
 %7Bthreshold%7D%5Ct%7Bun
@@ -3564,39 +3564,124 @@
 d%7D%5Ct
-%7Buniqs.shape%5B0%5D%7D%5Ct%7Bmax(counts)%7D
+Number of clusters: %7Buniqs.shape%5B0%5D%7D%5CtLargest cluster%7Bmax(counts)%7D. See the output file (%7Bjsonout%7D) for more details
 %22)%0A 

Clarifying output