from __future__ import division
from __future__ import print_function

from pathlib import Path
import sys
project_path = Path(__file__).resolve().parents[1]
sys.path.append(str(project_path))

import tensorflow as tf
import os
import scipy.sparse as sp
import numpy as np
from sklearn.metrics import euclidean_distances
import mygene
from collections import defaultdict
import csv
import markov_clustering as mc
import networkx as nx
import operator

# Settings
flags = tf.compat.v1.flags
FLAGS = flags.FLAGS
flags.DEFINE_string('dataset', 'brc_microarray_usa', 'Dataset string.')
flags.DEFINE_string('embeddings', 'embeddings.txt', 'Name of file that contains the latent node features.')
flags.DEFINE_string('embedding_method', 'gcn', 'Name of the embedding method.')

flags.DEFINE_string('clustering_method', 'mcl', 'Name of the clustering method.')
flags.DEFINE_string('clusters_output', 'clusters.txt', 'Name of desired output file of obtained clusters.')
flags.DEFINE_string('clusters_symbols', 'clusters.symbols.txt', 'Name of output file of ensembl names of cluster genes.')

#Check data availability
if not os.path.isfile("{}/data/output/{}/embedding/{}/{}".format(project_path, FLAGS.dataset, FLAGS.embedding_method, FLAGS.embeddings)):
    sys.exit("{} file is not available under /data/output/{}/embedding/{}/".format(FLAGS.embeddings, FLAGS.dataset, FLAGS.embedding_method))

if not os.path.isdir("{}/data/output/{}/clustering/{}/{}".format(project_path, FLAGS.dataset, FLAGS.clustering_method, FLAGS.embedding_method)):
    os.makedirs("{}/data/output/{}/clustering/{}/{}".format(project_path, FLAGS.dataset, FLAGS.clustering_method, FLAGS.embedding_method))

print("----------------------------------------")
print("----------------------------------------")
print("Clustering configuration:")
print("Dataset: {}".format(FLAGS.dataset))
print("Embedding method: {}".format(FLAGS.embedding_method))
print("Clustering Method: {}".format(FLAGS.clustering_method))
print("----------------------------------------")
print("----------------------------------------")
adj = sp.load_npz("{}/data/output/network/ppi.adjacency.npz".format(project_path))        

def save_clusters(clusters):
    with open("{}/data/output/{}/clustering/{}/{}/{}".format(project_path, FLAGS.dataset, FLAGS.clustering_method, FLAGS.embedding_method, FLAGS.clusters_output), "w", newline='', encoding="utf-8") as f:
        w_clusters = csv.writer(f, delimiter ='\t')
        for line in clusters:
            line = list(line)
            line.insert(0, "Center_{}")#Add an empty center to make the clutser row identical to the one generated by AP
            w_clusters.writerow(line)

#===============================================================================
# def prepare_clusters_ensembl_symbols(clusters_file="clusters.txt"):
#     
#     Ids = np.genfromtxt("{}/data/output/network/ppi.ids.txt".format(project_path), dtype=np.dtype(str), delimiter="\t")
#     protein_ids = {}
#     ids_protein = {}
#     for i in range(Ids.shape[0]):
#         protein_ids[Ids[i,0]]=Ids[i,1]
#         ids_protein[Ids[i,1]]=Ids[i,0]
#             
#     mg = mygene.MyGeneInfo()
#     map_ensp_symbol={}
#     print("Request gene symbols by gene query web service...")
#     annotations = mg.querymany(protein_ids.keys(), scopes='ensembl.protein', fields='symbol', species='human')
#     #For each query map ENSPs to the gene symbol
#     for response in annotations:
#         ensp = response['query']
#         if('symbol' in response):
#             map_ensp_symbol[ensp] = response['symbol']
#         else:
#             map_ensp_symbol[ensp] = ensp
#                     
#     Clusters = open("{}/data/output/{}/clustering/{}/{}/{}".format(project_path, FLAGS.dataset, FLAGS.clustering_method, FLAGS.embedding_method, FLAGS.clusters_output), encoding="utf-8")
#     with open("{}/data/output/{}/clustering/{}/{}/{}".format(project_path, FLAGS.dataset, FLAGS.clustering_method, FLAGS.embedding_method, FLAGS.clusters_symbols), "w", newline='', encoding="utf-8") as f:
#         w_cluster = csv.writer(f, delimiter ='\t')
#         for line in Clusters:
#             line = line.strip()
#             columns = line.split("\t")
#             cl = []
#             for i in range(1, len(columns)): #Skip first column that contains the exemplar
#                 if(ids_protein[columns[i]] in map_ensp_symbol.keys()):
#                     cl.append(map_ensp_symbol[ids_protein[columns[i]]])
#             w_cluster.writerow(cl)
#     Clusters.close()
#===============================================================================
    
#===============================================================================
# Read the node embeddings                
#===============================================================================
embeddings = np.genfromtxt("{}/data/output/{}/embedding/{}/{}".format(project_path, FLAGS.dataset, FLAGS.embedding_method, FLAGS.embeddings), dtype=np.float32)

#===============================================================================
# Start the clustering by geometric affinity propagation
#===============================================================================
print("Clustering by Markov clustering (MCL)...")
print("------------------------------")
similarities = -euclidean_distances(embeddings, squared=True)

network = nx.from_scipy_sparse_matrix(adj)
matrix = nx.to_scipy_sparse_matrix(network)
nnz_idx = matrix.nonzero()
for i,j in zip(*nnz_idx):
    matrix[i,j] = np.exp(similarities[i,j])

#===============================================================================
# print("Searching for best inflation parameter...")
# print("------------------------------")
# # perform clustering using different inflation values from 1.5 and 2.5
# # for each clustering run, calculate the modularity
# cluster_inflations = {}
# for inflation in [i / 10 for i in range(15, 26)]:
#     result = mc.run_mcl(matrix, inflation=inflation)
#     clusters = mc.get_clusters(result)
#     Q = mc.modularity(matrix=result, clusters=clusters)
#     cluster_inflations[inflation] = Q
#     print("inflation:", inflation, "modularity:", Q)
#  
# print("------------------------------")
# best_inflation = max(cluster_inflations.iteritems(), key=operator.itemgetter(1))[0]  
# print("Best inflation:", best_inflation, "Best modularity:", cluster_inflations[best_inflation])
# print("------------------------------")
#===============================================================================
best_inflation = 1.4
print("Clustering by Markov clustering (MCL) with inflation {}...".format(best_inflation))   
result = mc.run_mcl(matrix, inflation=best_inflation, verbose=True)           # run MCL with default parameters
clusters = mc.get_clusters(result)    # get clusters

n_clusters = len(clusters)
print("Number of obtained clusters: {}".format(n_clusters))
 
if n_clusters==0:
    sys.exit("Tune hyper-parameters for convergence")
  
print("Save clustering results ...")
save_clusters(clusters)
print("Clustering results saved")
 
#===============================================================================
# print("Preparing gene symbols for saved clusters ...")
# prepare_clusters_ensembl_symbols(clusters_file=FLAGS.clusters_output)
# print("Clusters with gene symbols saved")
#===============================================================================
