#!python
import random
from pkcm.utils import read_helices_in_file, read_pkcm_file
import tempfile
import argparse
import subprocess

parser = argparse.ArgumentParser(description="Emits sequences from a pseudoknotted covariance model.")

parser.add_argument('pkcm_file', type=str, help='.pkcm file to emit from')
parser.add_argument('-N', type=int, help="number of samples to get. (default=10)", default=10)

args = parser.parse_args()
N = args.N

helices = read_helices_in_file(args.pkcm_file)

transition_probabilities, in_ngbh, emission_probabilities, start_state = read_pkcm_file(open(args.pkcm_file).readlines())


def sample(list_of_outcomes):
    weight_list = [2**w for _,w in list_of_outcomes]
    cumulative_weight = 0
    cumulative_weights = []
    for k, p in enumerate(weight_list):
        cumulative_weight += p
        cumulative_weights.append(cumulative_weight)

    r = random.random()

    for k, cw in enumerate(cumulative_weights):
        if r-cw < 0:
            return list_of_outcomes[k][0]

    return list_of_outcomes[-1][0] # sometimes we can have
                                   # cumulative weights
                                   # not exactly summing to 1
                                   # because of precision errors

for _ in range(N):
    sequences = []
    for helix in helices:
        try:
            state = start_state[helix.name]
        except KeyError:
            no_in = [s for s in transition_probabilities[helix.name].keys() if s not in in_ngbh[helix.name].keys()]
            state = random.choice(no_in)
        left_seq = ""
        right_seq = ""
        while True:
            if state in emission_probabilities[helix.name].keys():
                outcome = sample(emission_probabilities[helix.name][state])
                if state.endswith('L'):
                    left_seq += outcome
                if state.endswith('R'):
                    right_seq = outcome + right_seq
                if state.endswith('P'):
                    left_seq += outcome[0]
                    right_seq = outcome[0] + right_seq
            if state in transition_probabilities[helix.name].keys():
                state = sample(transition_probabilities[helix.name][state])
            else:
                break
        sequences.append((helix.outer_left, left_seq))
        sequences.append((helix.inner_left, right_seq))

    print(''.join([s[1] for s in  sorted(sequences, key=lambda x : x[0])]))

