#!/usr/bin/env python3

"""Threading script

Usage:
  main.py -i <input> [-c <cutoff>] [-t <tplGap>] [-q <queryGap>]
          [-o <nbOutput>] [-b <bench>] -d <baseDir>

Options:
  -h --help                  help
  --version                  version of the script
  -i --input = query         input query file
  -c --cutoff = cut          the cutoff for the distance,
                             must be inferior to 15A [default: 15]
  -t --tplGap = tgap         penality for gap in the template
                             sequence [default: 0]
  -q --queryGap = qgap       penality for gap in the QUERY
                             sequence [default: 0]
  -o --nbOutput = nbO        number of outputs [default: 3]
  -b --bench = bench_mode    run the in 'benchmark mode'[default: False]
  -d --baseDir = base_dir    directory where we will fetch the different
                             data
"""

import os
import multiprocessing as mp
import time as t
import functools as ftools
import numpy as np
import pandas as pd
from docopt import docopt
import src.score as sc
import src.manage_io as mio
import src.rerank as rk
import src.checking_argument as cA
import src.parallelized as pll



# MAIN
if __name__ == "__main__":
    MAIN_TIME = t.time()
    cA.check_arguments(docopt(__doc__, version='0.1'))
    QUERY = os.path.basename(cA.QUERY).split('.')[0]  # Name of the QUERY
    print("  CURRENT QUERY = ", QUERY)

    # 0) Parsing of all needed files:
    with open(cA.QUERY, 'r') as input_foldrec_file: # query.foldrec
        DICO_ALIGNS, DICO_SS = mio.parse_foldrec(input_foldrec_file)

    # Get the list of families that are not suitable (several chains)
    # and delete them from the DICO_ALIGNS:
    with open(cA.BASE_DIR + "clean.log", 'r') as clean_log_file:
        NOT_TO_TAKE = clean_log_file.read().splitlines()
    for name in NOT_TO_TAKE:
        del DICO_ALIGNS[name]

    with open(cA.BASE_DIR + "Codes/Params/dope.par", 'r') as input_dope_file:
        DOPE_ARR_CA, INDEX_AA = mio.dope_to_arr_ca(input_dope_file)

    with open(cA.BASE_DIR + "Data/METAFOLD.list", 'r') as metafold_file:
        FILES_TO_TAKE = mio.parse_metafold_list(metafold_file)

    # Pandas DataFrame synthetizing all the scores:
    NB_ALN = len(DICO_ALIGNS.keys())
    SCORES_TYPES = ('alignment', 'co-evolution', 'SS', 'thread_CA', 'modeller')
    DF_SCORES = pd.DataFrame(data=np.zeros((NB_ALN, len(SCORES_TYPES)),
                                           dtype=float),
                             index=DICO_ALIGNS.keys(),
                             columns=SCORES_TYPES)
    for name in DICO_ALIGNS:
        # Add pdb_id for each template and a list of residus object which
        # contains their coordinates:
        DICO_ALIGNS[name].init_structure(FILES_TO_TAKE,
                                         cA.BASE_DIR + "Data/HOMSTRAD/" +
                                         name + '/')
        # Add alignment score to the df of scores
        DF_SCORES.loc[name, 'alignment'] = DICO_ALIGNS[name].score

    REF_TIME_SCORE = t.time()
    if not os.path.isdir("results"):
        os.mkdir("results")

    # 2) Co-evolution score:
    print("\t\t{}\n\t\t* Co-Evolution Score *\n\t\t{}".format("*"*22,
                                                              "*"*22))
    SET_INTERACT, LIST_INTERACT = sc.pre_proc_co_evol(QUERY)
    REF_COEVOL = t.time()
    # Calculates the co-evol score associated with each alignment:
    for fam_name in DICO_ALIGNS:
        ev_coevol = sc.CoEvolution(DICO_ALIGNS[fam_name])
        DF_SCORES.loc[fam_name, 'co-evolution'] = ev_coevol.co_evol(SET_INTERACT,
                                                                    LIST_INTERACT)
    print("Time Co-evolution = ", t.time() - REF_COEVOL)
    print("Co-Evolution Score calculation finished!\n")

    # 3) Secondary structures score:
    print("\t\t{}\n\t\t* 2 Structure Score *\n\t\t{}".format("*"*22, "*"*22))
    for fam_name in DICO_ALIGNS:
        SS_obj = DICO_SS[fam_name]
        SS_obj.ss_score() # Compute SS_score
        DF_SCORES.loc[fam_name, 'SS'] = SS_obj.score
    print("Secondary Structure Score calculation finished !\n")

    # 4) Parallelized threading with CA:
    print("\t\t{}\n\t\t* Threading Score CA *\n\t\t{}".format("*"*19, "*"*19))
    NB_CPU = mp.cpu_count() - 1
    print("Threading parallelized on", NB_CPU, "CPUs")
    MY_POOL = mp.Pool(NB_CPU)
    PARAL_THR_CA = ftools.partial(pll.parallelized_threading,
                                  dico_aligns=DICO_ALIGNS,
                                  dope_arr=DOPE_ARR_CA,
                                  index_aa=INDEX_AA)
    RES_THRD_CA = MY_POOL.map(PARAL_THR_CA, DICO_ALIGNS.keys())
    MY_POOL.close()

    # Add this new threading score the dataFrame of scores:
    for res in RES_THRD_CA:
        DF_SCORES.loc[res[0], "thread_CA"] = res[1]
    print("Threading CA score calculation finished !\n")

    # 6) Parallelized modeller:
    print("\t\t{}\n\t\t* Modeller Score *\n\t\t{}".format("*"*18, "*"*18))
    if not os.path.isdir("results/modeller_reslt"):
        os.makedirs("results/modeller_reslt")
    if os.path.isfile("problem_modeler.log"):
        os.remove("problem_modeler.log")
    MY_POOL = mp.Pool(mp.cpu_count()-1)
    MOD_FUNC = ftools.partial(pll.parallelized_modelling,
                              dico_aligns=DICO_ALIGNS)
    RESULTS_MOD = MY_POOL.map(MOD_FUNC, DICO_ALIGNS.keys())
    MY_POOL.close()
    print("Modeller Score calculation finished!")
    print("\n*****************************************")

    for res in RESULTS_MOD:
        DF_SCORES.loc[res[0], "modeller"] = res[1]

    print("Runtime for all scores = {}".format(t.time() - REF_TIME_SCORE))

    # Generation of the different required outputs:
    if not cA.BENCH_MODE:
        # Writing CA in PDB after threading:
        GLOB_SC_NAME = 'glob_bench'
        GLOBAL_SC_SORTED = rk.compute_global_sc(DF_SCORES, GLOB_SC_NAME,
                                                "results/bench/")

        for i in range(0, cA.NB_OUTPUT):
            top_hit = GLOBAL_SC_SORTED.loc[i, 'index']
            top_fam = DICO_ALIGNS[top_hit]
            # The aln needs to be cleaned for the function to work:
            top_fam.clean_algn()
            mio.write_pdb_out(top_fam)

        # Reranking families by final global score:
        rk.write_ranked_models(GLOBAL_SC_SORTED, GLOB_SC_NAME)
        print("Outputs written ! They are in the results/ directory\n")

    # NEEDED TO BENCHMARK:
    else:
        if not os.path.isdir("results/bench/"):
            os.makedirs("results/bench/")

        DF_SCORES.to_csv("results/bench/" + QUERY + ".csv")

    print("TOTAL_RUNTIME =", t.time() - MAIN_TIME, '\n')
