import os
import re
from os import path
from helperlibs.wrappers.io import TemporaryDirectory
from antismash import utils
from orderfinder import find_clusterpksnrpsgenes
from nrpspksdomainalign import substrates_specificity_phylo as ssp
from nrpspksdomainalign import index_domains as index_domains
from nrpspksdomainalign import get_msa_dist_matrix as mdm
from nrpspksdomainalign import calculate_distance_transATPKS as cdt
from nrpspksdomainalign import plot_domain_align as pda



def _get_transatpks_geneclusters(pksnrpsvars, seq_record):
    nrpspksclusters = list(set(utils.get_cluster_features_of_type(seq_record, "transatpks")))
    genes_in_cluster = {}
    for cluster in nrpspksclusters:
        cluster_id = utils.get_cluster_number(cluster)
        cluster_genes = [utils.get_gene_id(feature) for feature in find_clusterpksnrpsgenes(cluster, pksnrpsvars.pksnrpscoregenes)]
        genes_in_cluster[cluster_id] = cluster_genes

    return genes_in_cluster


def _get_nrpspks_domains_ks(pksnrpsvars, seq_record, domain):
    transatpks_geneclusters = _get_transatpks_geneclusters(pksnrpsvars, seq_record)
    transatpks_genes = list(set([g for g_list in transatpks_geneclusters.values() for g in g_list]))

    ksnames = []
    ksseqs = []
    if len(transatpks_geneclusters) >= 1:
        job_id = seq_record.id

        for feature in pksnrpsvars.pksnrpscoregenes:
            start_cds = str(feature.location.nofuzzy_start)
            end_cds = str(feature.location.nofuzzy_end)
            strand = feature.location.strand
            if strand == 1:
                strand_char = '+'
            else:
                strand_char = '-'
            loc = '-'.join((start_cds, end_cds))
            prot_id = product = ''
            if 'protein_id' in feature.qualifiers:
                prot_id = feature.qualifiers["protein_id"][0]
            if 'product' in feature.qualifiers:
                product = feature.qualifiers["product"][0].replace(' ', '_').replace('|', '')
                # We use | as a separator later
                assert '|' not in product, product

            gene_id = utils.get_gene_id(feature)

            if gene_id in transatpks_genes:

                domaindetails = pksnrpsvars.domaindict[gene_id]

                nr = 0
                for tab in domaindetails:
                    if tab[0] == domain:
                        nr += 1
                        start = int(tab[1])
                        end = int(tab[2])
                        loc_domain = '-'.join((str(start), str(end)))
                        ks_index = ''.join(('KS', str(nr)))

                        name1 = '|'.join(
                            [''.join(['>', job_id]), 'c', loc, strand_char, gene_id, product, prot_id, loc_domain, ks_index])
                        name = re.sub(r'(\:|\'|\(|\)|\,|\?|\;)', '', name1)
                        seq = str(utils.get_aa_sequence(feature))[start:end]

                        ksnames.append(name)
                        ksseqs.append(seq)
    return ksnames, ksseqs




def _index_domains_per_cluster(genes_per_cluster, domains_annotation, domain_ab):
    domains_per_cluster = {}
    for c in genes_per_cluster.keys():
        k = c
        v = {}
        gene_list = genes_per_cluster[c]
        for d in domains_annotation.keys():
            d_info = d.split("|")
            d_info[1] = d_info[1]+str(k)
            new_d = '|'.join(d_info)
            g = d.split('|')[4]
            if g in gene_list:
                v[new_d] = domains_annotation[d]

        v_domain_index = index_domains.run_index_domain(domain_id_list=v.keys(), domain_ab = domain_ab)
        v_new = {}
        for id_old in v.keys():
            id_new = v_domain_index[id_old]
            v_new[id_new] = v[id_old]


        domains_per_cluster[k] = v_new

    return domains_per_cluster




def classify_nrpspks_domains_ks(pksnrpsvars, seq_record, options):

    #todo: add an argument for the number of selected BGCs in run_antismash via parser.add_argument(),the parsed info will be store in options

    bgcs_nr = options.transatpks_da_cutoff

    with TemporaryDirectory(change=True):

        options.classify_domain_outputfolder_align = path.abspath(path.join(options.raw_predictions_outputfolder, "classified_domain", "align_ks"))
        if not os.path.exists(options.classify_domain_outputfolder_align):
            os.makedirs(options.classify_domain_outputfolder_align)

        options.classify_domain_outputfolder_tree = path.abspath(path.join(options.raw_predictions_outputfolder, "classified_domain", "tree_ks"))
        if not os.path.exists(options.classify_domain_outputfolder_tree):
            os.makedirs(options.classify_domain_outputfolder_tree)


        ks_names, ks_seqs = _get_nrpspks_domains_ks(pksnrpsvars, seq_record, domain="PKS_KS")   #todo: write ks_names and ks_seqs into a fasta file
        genes_per_transatpks = _get_transatpks_geneclusters(pksnrpsvars, seq_record)
        nrpspksdomainalign_dir = utils.get_full_path(__file__, "nrpspksdomainalign")
        data_dir = path.join(nrpspksdomainalign_dir, "data")
        annotation_out_dir = os.path.split(options.classify_domain_outputfolder_align)[0]

        genomeSize_info_seq_record = len(seq_record.seq)
        speciesName_seq_record = 'Query sequence'


        #align per novel ks to all training ks domains
        nr = 1
        for ksname, ksseq in zip(ks_names, ks_seqs):
            ssp.alignment_per_ks(seq=ksseq, name=ksname, data_dir=data_dir, align_dir=options.classify_domain_outputfolder_align, nr=nr)
            nr += 1

        #get annotation for KS_domains, a dict where key is domain id and value is the annotation
        KS_annotation = ssp.run_pipeline_perks_parallel(bgc_id=seq_record.id, out_dir=annotation_out_dir,
                                                        align_dir=options.classify_domain_outputfolder_align,
                                                        tree_dir=options.classify_domain_outputfolder_tree,
                                                        data_dir=data_dir, cutoff1=0.9, cutoff2=0.9, tree_conserve1=0.8,
                                                        tree_conserve2=0.8, cutoff1_new=0.7, cutoff2_new=0.7,
                                                        max_core=options.cpus)
        #todo: put all annotation of A-domains in A_annotation = ..., then merge A_annotation and KS_annotation into KSA_annotation= ...., and in the following function, index KS and A together in one assembly line


        #get indexed_KS_domain and annotated domains
        KS_per_cluster = _index_domains_per_cluster(genes_per_cluster=genes_per_transatpks, domains_annotation=KS_annotation, domain_ab = "KSA")
        # todo KSA_per_cluster = _index_domains_per_cluster(genes_per_cluster=genes_per_tranatpks, domains_annotation=KAS_annotation, domain_ab="KSA")

        # generate the txt file containing msa dist matrix for all domains
        KS_msa_dist = mdm.run_get_msa_dist_matrix(domain_names=ks_names, domain_seqs=ks_seqs, data_dir=data_dir,training_seq_filename="KS_rawseq_pred_training_transATPKS.txt", out_dir=annotation_out_dir, indexed_domain_per_cluster=KS_per_cluster)

        #get most similar assembly lines based on the calculated dist score
        #todo: modify the scipt so that it can extract the top
        similar_bgc_per_cluster, new_cluster, new_cluster_index = cdt.run_calculate_distance(data_dir = data_dir, seq_simScore = KS_msa_dist, ksa_per_new_cluster=KS_per_cluster, cutoff_bgc_nr = bgcs_nr)

        #generate figure output to html
        for b in similar_bgc_per_cluster.keys():
            b_id = re.sub("c",'', b.split('|')[1])
            dist_df = similar_bgc_per_cluster[b][1]
            annotateKS_training = os.path.join(data_dir, "annotateKS_perCompound_v4_uniPK.txt")
            annotateKS_novel = os.path.join(data_dir, "annotateKS_per_pred_transATPKS.txt")
            color_code_file = os.path.join(data_dir, "Clade_PhyloNode_v3_color.txt")
            KS_info_file_list = [os.path.join(data_dir, file) for file in ['KSinCompound_CladeInfo_v4_exclude.txt', 'KSinCompound_CladeInfo_pred_transATPKS.txt']]
            genomeSize_file_list = [os.path.join(data_dir, file) for file in['genome_size_training_transATPKS.txt', 'genome_size_pred_transATPKS.txt']]
            speciesName_file_list = [os.path.join(data_dir, file) for file in['species_ID_training_transATPKS.txt', 'species_ID_pred_transATPKS.txt']]
            outfile_svg_temp = os.path.join(annotation_out_dir, ''.join(["domain_align_cluster", b_id, ".svg"]))
            svg_fout = os.path.join(annotation_out_dir, ''.join(["domain_align_cluster", b_id, "_modify.svg"]))

            pda.run_plot_network_BGC_list(dist_df=dist_df,
                                  new_bgc_id=[b],
                                  new_cluster=new_cluster,
                                  new_cluster_index=new_cluster_index,
                                  annotateKS_training=annotateKS_training,
                                  annotateKS_novel=annotateKS_novel,
                                  speciesName_file_list=speciesName_file_list,
                                  dist=0.2,
                                  color_code_file=color_code_file,
                                  outfile=outfile_svg_temp,
                                  plt_width=40,
                                  KS_info_file_list=KS_info_file_list,
                                  genomeSize_file_list=genomeSize_file_list,
                                  ksa_per_new_cluster=KS_per_cluster,
                                  genomeSize_new_bgc=genomeSize_info_seq_record ,
                                  speciesName_new_bgc=speciesName_seq_record,
                                  seq_simScore = KS_msa_dist,
                                  svg_fout=svg_fout)


