#!/usr/bin/env python3
"""

"""

from scipy.cluster.hierarchy import linkage, fcluster
import numpy as np


def parse_table(lines):
    col_names = [l.strip() for l in lines[0].split("\t")[2:]]
    row_names = []
    frequencies = []

    for line in lines[1:]:
        tokens = line.split("\t")
        row_names.append((tokens[0], tokens[1]))
        frequencies.append([float(t) for t in tokens[2:]])

    return row_names, col_names, frequencies


if __name__ == "__main__":
    import sys
    from os import path
    sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
    from contact_calc.transformations import parse_contacts

    # Parse command line arguments
    import argparse as ap
    parser = ap.ArgumentParser(description=__doc__, formatter_class=ap.RawTextHelpFormatter)
    optional = parser._action_groups.pop()
    required = parser.add_argument_group('required arguments')
    parser._action_groups.append(optional)  # added this line

    required.add_argument('--table_input',
                          required=True,
                          type=ap.FileType('r'),
                          metavar='FILE',
                          help='A fingerprint table generated by get_contact_fingerprints.py')
    required.add_argument('--num_row_clusters',
                          required=False,
                          default=5,
                          type=int,
                          metavar='INT',
                          help='Max number of interaction clusters')
    required.add_argument('--num_col_clusters',
                          required=False,
                          default=5,
                          type=int,
                          metavar='INT',
                          help='Max number of structure clusters')
    required.add_argument('--pymol_output',
                          required=True,
                          metavar='FILE',
                          type=str,
                          help='The name of the output ')

    args = parser.parse_args()
    print("Parsing table from " + args.table_input.name)
    rows, cols, freqs = parse_table(args.table_input.readlines())

    freq_matrix = np.matrix(freqs)
    m, n = freq_matrix.shape
    print("Building linkage for clustering .. ")
    row_linkage = linkage(freq_matrix, method='ward')
    col_linkage = linkage(freq_matrix.T, method='ward')

    print("Splitting into clusters based on linkage .. ")
    row_cluster_assignments = fcluster(row_linkage, args.num_row_clusters, criterion='maxclust') - 1
    col_cluster_assignments = fcluster(col_linkage, args.num_col_clusters, criterion='maxclust') - 1
    row_clusters = list(set(row_cluster_assignments))
    col_clusters = list(set(col_cluster_assignments))

    # Compute cluster-matrix frequency
    print("Computing cluster-matrix frequencies .. ")
    cluster_freq_matrix = np.zeros((len(row_clusters), len(col_clusters)))
    for row_cluster in row_clusters:
        row_indices = [i for i in range(m) if row_cluster_assignments[i] == row_cluster]
        for col_cluster in col_clusters:
            col_indices = [i for i in range(n) if col_cluster_assignments[i] == col_cluster]

            ixgrid = np.ix_(row_indices, col_indices)
            cluster_freq_matrix[row_cluster, col_cluster] = \
                np.sum(freq_matrix[ixgrid]) / (len(row_indices) * len(col_indices))

    print(cluster_freq_matrix)

    print("Writing to " + args.pymol_output)
    if args.pymol_output:
        for cluster in col_clusters:
            cols_in_cluster = [cols[i] for i in range(n) if col_cluster_assignments[i] == cluster]
            fname = "%s_%04d.pml" % (args.pymol_output, cluster)

            with open(fname, 'w') as outfile:

                for col in cols_in_cluster:
                    outfile.write("#load(%s)\n" % col)

                outfile.write("\n")

                for row_cluster in row_clusters:
                    rows_in_cluster = [rows[i] for i in range(m) if row_cluster_assignments[i] == row_cluster]
                    object_name = "cluster_%d_freq_%.3f" % (row_cluster, cluster_freq_matrix[row_cluster, cluster])
                    for row in rows_in_cluster:
                        resi1 = "chain %s & resi %s & name CA+C1'" % (row[0].split(":")[0], row[0].split(":")[2])
                        resi2 = "chain %s & resi %s & name CA+C1'" % (row[1].split(":")[0], row[1].split(":")[2])
                        outfile.write("distance %s, %s, %s\n" % (object_name, resi1, resi2))
                outfile.write("\n")


__license__ = "Apache License 2.0"
__maintainer__ = "Rasmus Fonseca"
__email__ = "fonseca.rasmus@gmail.com"
