#!/usr/bin/env python3

########################################################################

# Usage:
# cat sample_in.txt | python3 fasta_stripper.py | python3 cons.py

import profile_matrix
import sys
import numpy as np

def cons(pm, regex_output=False):
    """Takes the output from profile_matrix and constructs
    a consensus string from it. Yeah!

    You can even make it output a regex in the case of
    multiple highest values, but otherwise it just defaults
    to the first one."""
    consensus = ''
    for row in pm[0].T:
        if isinstance(row.argmax(), np.ndarray):
            if regex_output:
                consensus += '['
                for x in row.argmax():
                    consensus = pm[1][int(x)]
                consensus += ']'
            else:
                consensus += row.argmax()[0]
        else:
            consensus += pm[1][int(row.argmax())]

    return consensus

if __name__=="__main__":
    raw = profile_matrix.profile_matrix(sys.stdin)
    print(cons(raw))

    s = ''
    for letter in raw[1]:
        s += letter
        s += ": "
        for n in raw[0][raw[1].index(letter)]:
            s += str(int(n))
            s += ' '
        print(s.strip())
        s = ''
