import sys
import os
import re
from itertools import islice
from sklearn.metrics.cluster import normalized_mutual_info_score

def get_true_clusters(true_taxon_file, level):
    # argument check
    if not (type(level) is int and level >=1 and level <=7):
        sys.stderr.write('invalid argument: level\n')
        sys.exit(1)

    true_clusters = {}  # cluster_label_id: [seq1, seq2, ...]
    keys = []   # save a keys list in case of keys order changing
    with open(true_taxon_file, 'r') as f:
        for line in f:
            content = line.strip().split('\t')
            seq_id = int(content[0])
            cluster_label = content[1].split(' ')[level - 1]
            if cluster_label in keys:
                true_clusters[ keys.index(cluster_label) ].append(seq_id)
            else:
                keys.append(cluster_label)
                true_clusters[ keys.index(cluster_label) ] = [seq_id]

    return true_clusters

def get_true_clusters_num(true_taxon_file, level):
    # argument check
    if not (type(level) is int and level >=1 and level <=7):
        sys.stderr.write('invalid argument: level\n')
        sys.exit(1)

    true_clusters = {}  # cluster_label_id: [seq1, seq2, ...]
    keys = []   # save a keys list in case of keys order changing
    with open(true_taxon_file, 'r') as f:
        for line in f:
            content = line.strip().split('\t')
            seq_id = int(content[0])
            cluster_label = content[1].split(' ')[level - 1]
            if cluster_label in keys:
                true_clusters[ keys.index(cluster_label) ].append(seq_id)
            else:
                keys.append(cluster_label)
                true_clusters[ keys.index(cluster_label) ] = [seq_id]

    return len(keys)

def get_clusters_from_qiime_otus(qiime_outs_file):
    my_clusters = {}    # cluster_label_id: [seq1, seq2, ...]
    cluster_cnt = 0
    with open(qiime_outs_file, 'r') as f:
        for line in f:
            content = line.strip().split('\t')
            my_clusters[cluster_cnt] = [int(x) for x in content[1:]]
            cluster_cnt += 1

    return my_clusters

def get_nmi(true_clusters, my_clusters):
    seqs = []   # keep a list of seq to make sure two assignment have same order
    true_assignment = []
    for key in true_clusters.keys():
        for seq_id in true_clusters[key]:
            seqs.append(seq_id)
            true_assignment.append(key)

    seq_to_cluster = {}
    for key in my_clusters.keys():
        for seq_id in my_clusters[key]:
            seq_to_cluster[seq_id] = key
    my_assignment = [seq_to_cluster[x] for x in seqs]

    return normalized_mutual_info_score(true_assignment, my_assignment)

def convert_single_est(est_file_1, est_file_2, est_file_3, cutoff, output_file):
    fa_id = [] 
    with open(est_file_1,'r') as f:
        while True:
            next_n = list(islice(f, 2))
            if not next_n:
                break
            header = next_n[0].strip()
            fa_id.append(header[1:])

    clean_map = []
    with open(est_file_2,'r') as f:
        for line in f:
            clean_map.append(line.strip().split())
        
    cluster = []
    with open(est_file_3,'r') as f:
        for line in f:
            content = line.strip().split('|')
            if float(content[0].strip()) == float(cutoff.strip()):
                for item in content[1:-1]:
                    cluster.append(item.split())


    with open(output_file,'w') as f:
        for i in range(len(cluster)):
            f.write('denovo' + str(i))
            c = cluster[i]
            for j in range(len(c)):
                merged_clique = clean_map[ int(c[j]) ]
                for k in range(len(merged_clique)):
                    f.write('\t' + fa_id[ int(merged_clique[k]) ])
            f.write('\n')


def convert_est_to_qiime_otus(cutoff):
    files = os.listdir('.')
    est_files = [x for x in files if re.match('epsrit_tree.*', x)]
    if len(est_files) % 8 != 0:
        sys.stderr.write('est file number error!\n')
    num_est = len(est_files)/8
    for i in range(num_est):
        convert_single_est('epsrit_tree_%d.fa' % i,
                           'epsrit_tree_%d_Clean.map' % i,
                           'epsrit_tree_%d_Clean.Clusters' % i,
                           cutoff, 'est_%d_otus.txt' % i)


def generate_combined_qiime_otus(original_fa):
    est_fa_2_origin_fa = {}
    cnt = 0
    with open(original_fa) as fa:
        while True:
            next_n = list(islice(fa, 2))
            if not next_n:
                break
            est_fa_2_origin_fa[cnt] = next_n[0].strip()[1:]
            cnt += 1

    files = os.listdir('.')
    cnt = 0
    with open('combined_otus.txt', 'w') as fo:
        for f in files:
            if re.match('est_\d+_otus.txt', f):
                with open(f, 'r') as fi:
                    for line in fi:
                        fo.write('denovo%d\t' % (cnt))
                        est_fa_items = line.strip().split('\t')[1:]
                        fo.write('%s\n' % '\t'.join([est_fa_2_origin_fa[int(x)] for x in est_fa_items]))
                        cnt += 1


if __name__ == '__main__':
    convert_est_to_qiime_otus(sys.argv[4])
    generate_combined_qiime_otus(sys.argv[1])
    print get_nmi(get_true_clusters(sys.argv[2], int(sys.argv[3])), get_clusters_from_qiime_otus('combined_otus.txt'))
