#!/usr/bin/python3
import warnings
import numpy as np
import subprocess
import argparse
from pkcm_rust_lib import rust_align_wrapper, QdbError, HmmQdbError
from pkcm.utils import get_sscons, constrained_helix_mapping_dict, extract_helix_state_graph, evaluate_alignment, best_possible_score_helix, read_pkcm_file, get_infernal_solution 
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
import sys
import time

parser = argparse.ArgumentParser(description="align a set of sequences in a fasta file to a pseudoknotted covariance model.")

parser.add_argument('pkcm_file', type=str, help='.pkcm file to align to')
parser.add_argument('td_file', type=str, help='.td file, nice tree decomposition of the graph of variables.')
parser.add_argument('fasta_file', type=str, help='.fasta of sequences to align')
parser.add_argument('output_file', type=str, help='output file')
parser.add_argument('seed_alignment', type=str, help='.seed seed alignment used to build the pkcm')
parser.add_argument('--beta', type=float, default=0.00001, help='parameter for Query-Dependent Banding, between 0 and 1.\
                    Higher is faster, but more approximate.')
parser.add_argument('--tau', type=float, default=0.001, help='parameter for Query-Dependent Banding, between 0 and 1.\
                    Higher is faster, but more approximate.')
parser.add_argument('--exact', action='store_true', help="Deactivates any heuristic")
parser.add_argument('--noqdb', action='store_true', help="Deactivates QDB")
parser.add_argument('--nohmm', action='store_true', help="Deactivates HMM")
parser.add_argument('--dense', action='store_true', help="Uses a denser contiguous array instead of a hashmap.\
                                                          Faster, but may use more memory.")
parser.add_argument('-N', type=int, default=sys.maxsize, dest='limit_num_sequences')

args = parser.parse_args()

outfile = open(args.output_file,'w', buffering=1)

# Getting the ss_cons and rf lines to include
# them in the output.
for line in open(args.seed_alignment).readlines():
    if line.startswith('#=GC SS_cons'):
        ss_cons = line.split(' ')[-1].rstrip('\n')
    if line.startswith('#=GC RF'):
        rf_line = line.split(' ')[-1].rstrip('\n')

# Re-figuring out the mapping between the 
# bounds of the helices and the 
# half helices <-> variables
partner = {}
half_helices = []
helices = []
for line in open(args.pkcm_file).readlines():
    if line.startswith("## END HELIX ARRANGEMENT DESCRIPTION"):
        break
    if line.startswith("Helix"):
        name = line.split(' ')[0]
        i = int(line.split(' ')[-4])-1
        j = int(line.split(' ')[-3])-1
        k = int(line.split(' ')[-2])-1
        l = int(line.split(' ')[-1].rstrip('\n'))-1

        half_helices.append((i,j))
        half_helices.append((k,l))
        helices.append((name,(i,j),(k,l)))

        partner[(k,l)] = (i,j)

half_helices = sorted(half_helices, key=lambda x : x[0])

cnt1 = 1
cnt2 = 2

import copy

pos = {}

for k in range(len(half_helices)):

    i,j = half_helices[k]
    pos[(i,j)] = (copy.copy(cnt1), copy.copy(cnt2))

    cnt1 += 1
    cnt2 += 1

# end of figuring out mapping


cons_column_writes = [] # list of dicts: column -> symbol. should be exactly one symbol for every column
insert_writes = []      # list of dicts: column -> the insert written between column and column+1
num_inserts = {}        # column: length of longest insert between column and column+1
ids = []

num_cons_cols = 0
for k, c in enumerate(list(rf_line)):
    if c != '.':
        num_cons_cols += 1

## retrieve HMM model
#print("READING HMM")
#start_reading = False
#for line in open(args.pkcm_file).readlines():
#    if line.startswith("## FULL HMM"):
#        start_reading = True
#        continue
#    if start_reading:
#        if line.startswith("transition probabilities"):
#            transition_probabilities_hmm = parse_transition_probas_hmm(line)
#        if line.startswith("match_emissions probabilities"):
#            match_emissions_probabilities_hmm = parse_transition_probas_hmm(line)
#        if line.startswith("insert_emissions probabilities"):
#            insert_emissions_probabilities_hmm = parse_transition_probas_hmm(line)


# map id -> seq
id2seq = {}

# loop over input sequences
cnt = 0
times = []
is_too_long_for_beta = []
for record in SeqIO.parse(args.fasta_file, "fasta"):
    if cnt >= args.limit_num_sequences:
        break
    cnt += 1

    ids.append(record.id)
    id2seq[record.id] = str(record.seq)

    print("aligning",record.id)
    outfile.write(" ".join(["aligning", record.id,'\n']))
    t0 = time.time()
    try:
        opt, mapping, state_lists = rust_align_wrapper(args.pkcm_file, 
                                                   args.td_file, 
                                                   str(record.seq),
                                                   args.beta,
                                                   args.tau,
                                                   args.exact,
                                                   args.dense,
                                                   args.noqdb,
                                                   args.nohmm)
    except QdbError as e:
        print(str(e))
        is_too_long_for_beta.append(record.id)
        continue

    except HmmQdbError as e:
        print(str(e))
        is_too_long_for_beta.append(record.id)
        continue

    times.append(time.time()-t0)

    # corrected score using NULL random model
    score = opt+2*len(str(record.seq))

    print("computation time ",time.time()-t0)
    print("id, opt, score", record.id, opt, score)
    print("mapping ", mapping)
    outfile.write(f"id {record.id} score {score} computation_time {time.time()-t0}\n")

    # null3 proba
    nA = float(record.seq.count('A'))
    nU = float(record.seq.count('U'))
    nG = float(record.seq.count('G'))
    nC = float(record.seq.count('C'))
    nTotal = nA+nU+nG+nC

    null3_log_prob = 0.0
    for c in ['A','U','G','C']:
        n = float(record.seq.count(c))
        if n > 0:
            null3_log_prob += n*np.log2((n/nTotal))

    seq_list = list(str(record.seq))

    # converting to dict
    map_dict = {}
    for key,val in mapping:
        map_dict[key] = val

    cons_column_write = {}
    insert_write = {}
    for col_index in range(num_cons_cols+2):
        insert_write[col_index] = [""]

    for helix_index in range(len(helices)):

        i,j = helices[helix_index][1]
        k,l = helices[helix_index][2]

        pos_i,pos_j = pos[(i,j)]
        pos_k,pos_l = pos[(k,l)]

        start_L,end_L = map_dict[pos_i]+1, map_dict[pos_j]
        start_R,end_R = map_dict[pos_k]+1, map_dict[pos_l]

        var_L = start_L-1
        var_R = end_R-1

        # shifting
        i += 1
        l += 1

        #print("starting indices",i,l)
        #print("map_dict", map_dict)
        #print(state_lists[helix_index])
        #print("left seq", "".join(seq_list[start_L-1:end_L]),"right seq","".join(seq_list[start_R-1:end_R]))

        def left_insert_pos(state, num_cons_cols):
            if state.startswith("ROOT"):
                return 0
            try:
                i = int(state.split('_')[1])
                if state.startswith("BEGR"):
                    return i-1
            except:
                assert(state.split('_')[1]=='-')
                i = 0
            return i

        def right_insert_pos(state, num_cons_cols):
            if state.startswith("ROOT"):
                return num_cons_cols
            try:
                i = int(state.split('_')[2])-1
            except:
                assert(state.split('_')[2]=='-')
                i = num_cons_cols
            return i

        for state in state_lists[helix_index]:
        # in what follows: i and l track the cons positions
        # var_L and var_R track the sequence positions
            if state.endswith("IL"):
                insert_write[left_insert_pos(state, num_cons_cols)].append(seq_list[var_L])
                var_L += 1
                continue

            if state.endswith("IR"):
                insert_write[right_insert_pos(state, num_cons_cols)] = [seq_list[var_R]] + insert_write[right_insert_pos(state, num_cons_cols)]
                var_R -= 1
                continue

            if state.startswith("MATP"):
                if state.endswith("MP"):
                    cons_column_write[i] = seq_list[var_L] 
                    cons_column_write[l] = seq_list[var_R]

                    var_L += 1
                    var_R -= 1

                if state.endswith("ML"):
                    cons_column_write[i] = seq_list[var_L] 
                    cons_column_write[l] = "-"

                    # no var_R decrement
                    var_L += 1

                if state.endswith("MR"):
                    cons_column_write[i] = "-" 
                    cons_column_write[l] =  seq_list[var_R]

                    var_R -= 1

                if state.endswith("D"):
                    cons_column_write[i] = "-" 
                    cons_column_write[l] = "-"
                    
                # for matp node: in any case we write both left/right
                # consensus columns
                i += 1
                l -= 1

            if state.startswith("MATL"):
                if state.endswith("ML"):
                    cons_column_write[i] = seq_list[var_L] 
                    var_L += 1

                if state.endswith("D"):
                    cons_column_write[i] = "-" 

                i += 1

            if state.startswith("MATR"):
                if state.endswith("MR"):
                    cons_column_write[l] = seq_list[var_R] 
                    var_R -= 1

                if state.endswith("D"):
                    cons_column_write[l] = "-" 

                l -= 1
                
    cons_column_writes.append(cons_column_write)
    insert_writes.append(insert_write)

for k in range(num_cons_cols+2):
    max_ins = 0
    for insert_write in insert_writes:
        max_ins = max(max_ins,len(insert_write[k])-1) # -1 because there is ""

    num_inserts[k] = max_ins


max_id_len = 0
for identifier in ids+['ss_cons','rf_line']:
    max_id_len = max(len(identifier)+1, max_id_len)

def pad_with_spaces(s, L):
    while len(s) < L:
        s += ' '
    return s

# TODO: right-inserts should be padded on the left, not
# the right
def pad_with_gaps(s, L):
    while len(s) < L:
        s += '-'
    return s

def pad_with_dots(s, L):
    while len(s) < L:
        s += '.'
    return s

def print_aligned_seq(identifier, num_cons_cols, insert_write, cons_column_write, num_inserts, max_id_len):

#    print("cons   write", sorted(cons_column_write.items(),key = lambda x : x[0]))
#    print("insert write", insert_write)

    # also storing the seq for validation purposes
    seq = ""

    print(pad_with_spaces(identifier, max_id_len),end='')
    outfile.write(pad_with_spaces(identifier, max_id_len))
    for col_index in range(num_cons_cols+2):
        if col_index > 0 and col_index <= num_cons_cols:
            print(cons_column_write[col_index], end='')
            outfile.write(cons_column_write[col_index])
            seq += cons_column_write[col_index]

        print(pad_with_gaps(''.join(insert_write[col_index]), num_inserts[col_index]),end='')
        outfile.write(pad_with_gaps(''.join(insert_write[col_index]), num_inserts[col_index]))
        seq += ''.join(insert_write[col_index])

    print("")
    outfile.write('\n')
    return seq

def print_aligned_cons(identifier, num_cons_cols, insert_write, cons_column_write, num_inserts, max_id_len):
    print(pad_with_spaces(identifier, max_id_len),end='')
    outfile.write(pad_with_spaces(identifier, max_id_len))
    for col_index in range(num_cons_cols+2):
        if col_index > 0 and col_index <= num_cons_cols:
            print(cons_column_write[col_index], end='')
            outfile.write(cons_column_write[col_index])

        print(pad_with_dots(''.join(insert_write[col_index]), num_inserts[col_index]),end='')
        outfile.write(pad_with_dots(''.join(insert_write[col_index]), num_inserts[col_index]))

    print("")
    outfile.write('\n')

# making objects with good format for ss_cons and gc_rf
ss_cons_col_write = {}
ss_cons_insert_write = {}
rf_line_col_write = {}
rf_line_insert_write = {}
curr_insert = ""
tracker = 0

for col_index in range(num_cons_cols+2):
    ss_cons_insert_write[col_index] = ""
    rf_line_insert_write[col_index] = ""

col_index = 0
for k,c in enumerate(list(rf_line)):
    if c!='.':
        ss_cons_col_write[col_index+1] = ss_cons[k]
        rf_line_col_write[col_index+1] = rf_line[k]
        col_index += 1
    
print("average single-sequence alignment time ", np.mean(times),"seconds")
outfile.write(" ".join(["average single-sequence alignment time ", str(np.mean(times)),"seconds\n"]))

# final print
print_aligned_cons("ss_cons", num_cons_cols, ss_cons_insert_write, ss_cons_col_write, num_inserts, max_id_len)
print_aligned_cons("rf_line", num_cons_cols, rf_line_insert_write, rf_line_col_write, num_inserts, max_id_len)

ids = [identifier for identifier in ids if identifier not in is_too_long_for_beta]

for num,identifier in enumerate(ids):
    reconstructed_seq = print_aligned_seq(identifier, num_cons_cols, insert_writes[num], cons_column_writes[num], num_inserts, max_id_len)
    reconstructed_seq = reconstructed_seq.replace('-','')
    #print("for id", identifier, "reconstructed seq is ", reconstructed_seq)
    #print("            whereas seq is ", id2seq[identifier])
    assert(reconstructed_seq==id2seq[identifier])
