#!/usr/bin/python3
from pkcm.utils import Node, Helix, adhoc_helix_decomposition, get_sscons, preambule_stk_file, alignment_lines_stk_file, filter_alignment, homemade_cmbuild, read_pkcm_lines_infernal_format, equalise_whitespace
import subprocess
import tempfile
import sys
import argparse

# parser configuration
parser = argparse.ArgumentParser(description='Builds a pseudoknotted covariance model (.pkcm file) from a multiple sequence alignment. Also outputs a .helix file specifying how the consensus structure was decomposed into helices.')

# main arguments
parser.add_argument('output_pkcm_filename', type=str, help='name of the file in which to write the pseudoknotted covariance model.')
parser.add_argument('inputfile_msa_name', type=str, help='name of the input multiple sequence alignment')

# optional arguments
parser.add_argument('--dump-helix-alignments', 
                    type=str, 
                    help="optionally, directory in which to dump alignment files of each helix.", 
                    dest='dump_helix_alignments',
                    metavar='DIR',
                    required=False)

parser.add_argument('--ignore-pk',
                    help="Ignore the pk helices (those with Aa,Bb,Cc...)",
                    dest='ignore_pk',
                    action='store_true',
                    required=False)

parser.add_argument('--use-infernal-format',
                    help="If present, use infernal (cmbuild) for each helix. Instead of simple Laplace smoothing and ``home-made'' format.",
                    dest='use_infernal',
                    action='store_true',
                    required=False)

# parsing args
args = parser.parse_args()

msa_filename = args.inputfile_msa_name
cm_outname = args.output_pkcm_filename
try:
    helix_msa_dir = args.dump_helix_alignments
except AttributeError:
    helix_msa_dir = None


# getting ss_cons
ss_cons, rf_line = get_sscons(open(msa_filename).readlines(), ignore_pk=args.ignore_pk)
rf_line = rf_line.replace('.','')
N = len(ss_cons)

# helix decomposition
helix_list = adhoc_helix_decomposition(ss_cons)

pkcm_file = open(cm_outname,'w')

if args.use_infernal:
    print('#USING_CMBUILD_BACKEND=true',file=pkcm_file)
    ss_cons_spacer = '<<<.>>>'
    gc_rf_spacer = '..GAC..'
else:
    print('#USING_CMBUILD_BACKEND=false',file=pkcm_file)
    ss_cons_spacer = '||'
    gc_rf_spacer = '||'

# storing helix decomposition description
print("## BEGIN HELIX ARRANGEMENT DESCRIPTION ##", file=pkcm_file)
for helix in helix_list:
    print(helix.name, helix.parenthesis_system, helix.print_boundaries(), file=pkcm_file)
    for node in helix.nodes:
        print(node.type, node.left_end, node.right_end,file=pkcm_file)
    print('-'*10,file=pkcm_file)
print("## END HELIX ARRANGEMENT DESCRIPTION ##", file=pkcm_file)

preambule = preambule_stk_file(open(msa_filename).readlines())
alignment_lines = alignment_lines_stk_file(open(msa_filename).readlines(),ignore_pk=args.ignore_pk)

all_lines = []

cm_lines_per_helix = {}

if not args.use_infernal:
    print("## BEGIN PKCM ##", file=pkcm_file)

for helix in helix_list:

    # isolating the part of the alignment concerning the helix
    alignment_file = tempfile.NamedTemporaryFile()
    filter_alignment(preambule, alignment_lines, helix, alignment_file,ss_cons_spacer=ss_cons_spacer,gc_rf_spacer=gc_rf_spacer) 
    alignment_file.seek(0)
    if helix_msa_dir is not None:
        helix_msa_file = open(helix_msa_dir+msa_filename.split('.')[0].split('/')[-1]+'_'+helix.name+'.seed','w')
        for line in alignment_file.readlines():
            helix_msa_file.write(line.decode('utf-8'))
        alignment_file.seek(0)

    if args.use_infernal:
        # file to store the resulting covariance model for that helix
        helix_cm_file = tempfile.NamedTemporaryFile()

        print("input given to cmbuild:")
        for line in alignment_file.readlines():
            print(line)
        alignment_file.seek(0)

        # infernal run
        cmd_args = ['cmbuild','--iins','--iflank','--hand','-F',helix_cm_file.name,alignment_file.name]
        subprocess.run(cmd_args, check=True)

        # dumping into overall pkcm file
        helix_cm_file.seek(0)
        cm_lines_per_helix[helix.name] = list(helix_cm_file.readlines())

        all_lines.append('## CM FILE for '+helix.name)
        for line in cm_lines_per_helix[helix.name]:
            print(line)
            all_lines.append(line.decode("utf-8"))

    else:
        print('-----------',file=pkcm_file)
        print(helix.name, helix.parenthesis_system, helix.print_boundaries(), file=pkcm_file)

        transition_edge_list, emission_list = homemade_cmbuild(alignment_file, helix)

        for state1, state2, weight in transition_edge_list:
            print('transition',state1,'->',state2,weight,file=pkcm_file)

        for state, nucl, weight in emission_list:
            print('emission',state,'->', nucl, weight, file=pkcm_file)

if not args.use_infernal:
    print("## END PKCM ##", file=pkcm_file)

if args.use_infernal:
    transition_probabilities, emission_probabilities = read_pkcm_lines_infernal_format(all_lines, helix_list)
    print("## BEGIN PKCM ##", file=pkcm_file)
    for helix in helix_list:
        print('-----------',file=pkcm_file)
        print(helix.name, helix.parenthesis_system, helix.print_boundaries(), file=pkcm_file)
        for key, value in transition_probabilities[helix.name].items():
            print('transition',key,value,file=pkcm_file)
        for state, value in emission_probabilities[helix.name].items():
            for nucl, weight in value.items():
                print('emission',state,'->',nucl,weight,file=pkcm_file)
    print("## END PKCM ##", file=pkcm_file)
    print("## BEGIN INDIVIDUAL CM FILES", file=pkcm_file)
    for helix in helix_list:
        for line in cm_lines_per_helix[helix.name]:
            print(line.decode("utf-8"),end='',file=pkcm_file)
    print("## END INDIVIDUAL CM FILES", file=pkcm_file)
else:
    print("## END PKCM ##", file=pkcm_file)

rf_line, ss_cons = equalise_whitespace(rf_line, '#=GC SS_cons '+ss_cons)

import pyhmmer
from pyhmmer import easel, plan7

alphabet = easel.Alphabet.rna()

with easel.MSAFile(args.inputfile_msa_name, digital=False) as msa_file:
    msa = next(msa_file)

digital_msa = msa.digitize(alphabet)

builder = plan7.Builder(alphabet)
background = plan7.Background(alphabet)

hmm, profile, optimized_profile = builder.build_msa(
    digital_msa, background
)

print("## FULL HMM ", file=pkcm_file)
print("transition probabilities", hmm.transition_probabilities, file=pkcm_file)
print("match_emission probabilities", hmm.match_emissions, file=pkcm_file)
print("insert_emission probabilities", hmm.insert_emissions, file=pkcm_file)
print("## END FULL HMM", file=pkcm_file)


pkcm_file.write('### CONSENSUS INFO'+'\n')
pkcm_file.write(ss_cons+'\n')
pkcm_file.write(rf_line+'\n')
