import multiprocessing
import os
import time
from itertools import combinations
from pathlib import Path
import numpy as np
import pandas as pd
from Comparison_digital_profiles_human import compare_pair
from Signal_digitalisation_human import MatricesExtractor
from utils import InputFileManager
from utils import create_dir_if_not_exist

np.random.seed(10)

add_small_random_value_to_random_comparison = True

num_cores = multiprocessing.cpu_count()
num_task = 5

# I/O directories
input_dir = os.path.join(os.getcwd(),"/home/gg486/rds/rds-yhbl2-genehunter/bics/Personal_folders/gg486/Riboseq_Human_control/Serial/python/Esperimento_random/Selected_genes/LC509_no_soglia") #"/home/gg486/rds/rds-yhbl2-genehunter/bics/Personal_folders/gg486/Riboseq_Human_control/Serial/python/Cancer_bedfile_original") #/rds/project/yhbl2/rds-yhbl2-genehunter/bics/Personal_folders/gg486/Riboseq_Human_control/alignment_coding_sequences/alignToCDS_bed") #bedfiles_cancer_pairwise/") #bed_files_475_control/")  # get the path to the data input directory

match_scores_output_dir = os.path.join(os.getcwd(),"matrix_python_LC509_nosoglia_27") #"matrix_python_CONTROL_all") #"matrix_python_472/match_scores_472/")  # Sets the directory where all the saved outputs will be stored
reproducible_sequence_output_dir = os.path.join(os.getcwd(), "matrix_python_LC509_nosoglia_27/reproducible_sequences/") #matrix_python_CONTROL_all/reproducible_sequence/")#"matrix_python_472/reproducible_sequence/")  # Sets the directory where all the saved outputs will be stored
genes_lengths_path = os.path.join(os.getcwd(), "../../CDShumanGenesLengths.txt")  # path to upload the file containing each gene's ID and the correspondent gene length
histogram_plot_path = os.path.join(os.getcwd(), "genes_histograms_LC509_nosoglia_27/")  # path to upload the file containing each gene's ID and the correspondent gene length
intermediate_results = os.path.join(os.getcwd(), "intermediate_results_LC509_nosoglia_27/")
random_comparisons_folder_match_scores = os.path.join(os.getcwd(), "random_comparisons_folder_match_scores_LC509_nosoglia/")
random_comparisons_folder_matrix_01 = os.path.join(os.getcwd(), "random_comparisons_folder_matrix_01_LC509_nosoglia/")

create_dir_if_not_exist([input_dir, match_scores_output_dir, histogram_plot_path, reproducible_sequence_output_dir, intermediate_results])

num_comparison = 200  #200 e 50 per fare 10000  # NOTA: numero di confronti random da eseguire per ogni coppia di file bed

save_random_match_scores = False
save_matrix_01 = False


def extract_gene_list(genes, bed_files_dicts):
    gene_lists = []
    for bed_files_dict in bed_files_dicts:
        bed_file = bed_files_dict["bed_file"]
        table_FP = bed_file["Chromosome"].value_counts().sort_index().rename_axis('GeneID').reset_index(name='ReadsCounts')
        table_FP_Geneslengths = pd.merge(table_FP, genes, on="GeneID")

        bed_table_FP_reduced = bed_file.join(table_FP_Geneslengths.set_index('GeneID'), on='Chromosome').set_index("Chromosome")
        gene_name_list = []
        for name_of_the_group, group in bed_table_FP_reduced.groupby("Chromosome"):
            gene_name_list.append(name_of_the_group)
        gene_lists.append(gene_name_list)

    # gets the gene list for each matrix
    gene_lists = [pd.DataFrame(f, columns={"GeneID"}) for f in gene_lists]
    # gets the genes in common for each matrix
    gene_list = gene_lists[0]
    if len(gene_lists) > 1:
        for i in range(1, len(gene_lists)):
            l = gene_lists[i]
            gene_list = gene_list.merge(l, on="GeneID")

    gene_list = gene_list.to_numpy().squeeze()
    return gene_list

def random_comparison(arguments, rep_num=0):  # esegue i confronti random
    matrices_extractors, genes, gene_list = arguments
    matrix_01_pair = []
    for matrices_extractor in matrices_extractors:
        bed_file_name = matrices_extractor["bed_file_name"]
        me = matrices_extractor["me"]
        # extract the matrices
        # print("random comparison num: " + str(rep_num) + " extract 01 and coverage from: " + str(bed_file_name))
        # start = time.time()
        matrix01, coverage = me.extract_matrices(areReadsRandomized=True, add_small_random_value=add_small_random_value_to_random_comparison, rep_num=rep_num)
        matrix_01_pair.append({"matrix":matrix01, "file_name": bed_file_name})

        # print ("end in " + str(time.time() - start) + " sec(s)")
        if save_matrix_01:
            create_dir_if_not_exist([random_comparisons_folder_matrix_01])
            save_dir = os.path.join(random_comparisons_folder_matrix_01, bed_file_name)
            create_dir_if_not_exist([save_dir])
            matrix_01.to_csv(os.path.join(save_dir, str(rep_num) + ".csv"), index=True, header=True, decimal='.', sep=',', float_format='%.6f')
    # print(" start compare pairs ")
    # start = time.time()
    match_score, pair_names = compare_pair(matrix_01_pair, genes.set_index('GeneID'), gene_list)  # compara coppie di file bed  (di matrici 0-1)

    pair_names = Path(pair_names[0]).stem + ":" + Path(pair_names[1]).stem
    # print("end_comparison in " + str(time.time() - start) + "seconds")
    return {'pair_name': pair_names, 'match_score': match_score}


def compare_pair_n_times(bed_files_pair, genes, gene_list, n):
    # extract a pair of bed files
    match_scores = []

    matrices_extractors = []
    for bed_files_dict in bed_files_pair:  # FIX ME creazione delle classi estrattori a monte, verificare correttezza
        bed_file = bed_files_dict["bed_file"]
        bed_file_name = bed_files_dict["bed_file_name"]
        me = MatricesExtractor(bed_file, genes)
        matrices_extractors.append({"me": me, "bed_file_name": bed_file_name})

    arguments = matrices_extractors, genes, gene_list
    pool = multiprocessing.Pool(processes=num_task)
    res = []
    for i in range(n):
        res.append(pool.apply_async(random_comparison, [arguments]))
    pool.close()
    pool.join()
    for i in res:
        match_scores.append(i.get())

    if save_random_match_scores:
        create_dir_if_not_exist([random_comparisons_folder_match_scores])
        for i in range(len(match_scores)):
            save_dir = os.path.join(random_comparisons_folder_match_scores, match_scores[i]["pair_name"])
            create_dir_if_not_exist([save_dir])
            match_scores[i]["match_score"].to_csv(os.path.join(save_dir, str(i) + ".csv"), index=True, header=True, decimal='.', sep=',', float_format='%.6f')
    return match_scores


def compare_pair_n_times_serial(bed_files_pair, genes, gene_list, n):  # per ogni coppia di file bed esegue n confronti
    # extract a pair of bed files
    match_scores = []

    matrices_extractors = []
    for bed_files_dict in bed_files_pair:
        bed_file = bed_files_dict["bed_file"]
        bed_file_name = bed_files_dict["bed_file_name"]
        me = MatricesExtractor(bed_file, genes, bed_file_name)
        matrices_extractors.append({"me": me, "bed_file_name": bed_file_name})

    arguments = matrices_extractors, genes, gene_list

    for i in range(n):
        match_scores.append(random_comparison(arguments, i))  # effettua n confronti random

    if save_random_match_scores:
        create_dir_if_not_exist([random_comparisons_folder_match_scores])
        for i in range(len(match_scores)):
            save_dir = os.path.join(random_comparisons_folder_match_scores, match_scores[i]["pair_name"])
            create_dir_if_not_exist([save_dir])
            match_scores[i]["match_score"].to_csv(os.path.join(save_dir, str(i) + ".csv"), index=True, header=True, decimal='.', sep=',', float_format='%.6f')

    return match_scores


def main(num):
    print("num of core available: " + str(num_cores) + " used: " + str(num_task))

    ifm = InputFileManager(genes_lengths_path, input_dir)
    genes = ifm.get_genes_human()
    bed_files_dicts = ifm.get_bed_files()
    gene_list = extract_gene_list(genes, bed_files_dicts)

    # create pairs of bed files
    bed_files_pairs = [list(f) for f in combinations(bed_files_dicts, 2)]
    print("start fake matrix comparisons...")

    match_scores_list = []

    for bed_files_pair in bed_files_pairs:
        start = time.time()
        match_scores_list.append(compare_pair_n_times(bed_files_pair, genes, gene_list, num_comparison))
        end = time.time()
        print("fake matrix comparisons " + bed_files_pair[0]["bed_file_name"] + " vs " + bed_files_pair[1]["bed_file_name"] + " completed in " + str(end - start) + " sec(s)")

    #for bed_files_pair in bed_files_pairs:
    #    start = time.time()
    #    print ("start comparison: " + str(bed_files_pair[0]['bed_file_name']) + "vs" + str(bed_files_pair[1]['bed_file_name']))
    #    match_scores_list.append(compare_pair_n_times_serial(bed_files_pair, genes, gene_list, num_comparison))
    #    end = time.time()
    #    print("fake matrix comparisons " + bed_files_pair[0]["bed_file_name"] + " vs " + bed_files_pair[1]["bed_file_name"] + " completed in " + str(end - start) + " sec(s)")

    with open(intermediate_results + '/match_scores_list_' + str(num) + '_aggregating_' + str(num_comparison) + '_RANDOM_LC509_nosoglia_comparisons.npy', 'wb') as f:
        np.save(f, match_scores_list)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='riboseq reproducibility')
    parser.add_argument('--num', metavar='n', required=False, help='experiment identifier')
    args = parser.parse_args()
    main(num=args.num)
