import numpy as np
from itertools import count

def remove_linebreak(path):
    # path = 'clustalo-500_ph_new.fasta'
    stem = path.split('.')[0]
    ext = path.split('.')[1]
    input_data = path
    output_data = stem + '_stripped.' + ext

    with open(input_data, 'r') as in_file, \
         open(output_data, 'w') as out_file:
        
        for ii,row in enumerate(in_file):
            row = row.strip('\n')

            if ii > 0 and ">" in row:
                out_file.write('\n')

            out_file.write(row)

            if ">" in row:
                out_file.write('\n')

    return output_data

def remove_mafft_train(path):
    stem = path.split('.')[0]
    ext = path.split('.')[1]
    input_data = path
    output_data = stem + '_predict_chunk.' + ext

    with open(input_data, 'r') as fas_file,\
         open(output_data, 'w') as out_file:
        read_iter = iter(fas_file.readline, '')
        for fas_head, fas_data in zip(read_iter, read_iter):
            if ">TEST" in fas_head:
                # print(fas_head)
                out_file.write(fas_head.strip('\n\r')+'\n')
                out_file.write(fas_data.strip('\n\r')+'\n')

    return output_data

def change_to_sparse_categorical(symbols):
    canonical_symbols = '-ACDEFGHIKLMNPQRSTVWY'
    canonical_symbols = dict(zip(canonical_symbols, count()))
    # canonical_symbols = {'-': 0, 'A': 1, 'C': 2, 'D': 3, ...}
    new_symbols = []
    for s in symbols:
        if s in canonical_symbols:
            new_symbols.append(canonical_symbols[s])
        else:
            new_symbols.append(0)

    return np.array(new_symbols)

def fasta_to_numpy(mfile):
    data = []
    with open(mfile, 'r') as fas_file:
        read_iter = iter(fas_file.readline, '')
        for fas_head, fas_data in zip(read_iter, read_iter):
            fas_head = fas_head.strip('\n\r')
            fas_data = fas_data.strip('\n\r')
            d = change_to_sparse_categorical(list(fas_data))
            data.append(d)

    data = np.stack(data, 0)
    return data

def read_embedding_matrix(path):
    with open(path, 'r') as csv_file:
        c = None
        embedding_matrix = None
        for ii,row in enumerate(csv_file):
            if not row:
                continue
            if ii == 0:
                c = float(row.strip(',\r\n'))
                continue

            values = row.strip(',\r\n').split(',')
            values = [float(v) for v in values]
            embedding = np.array(values[1:])[None,...]

            if ii == 1:                
                embedding_matrix = embedding
            else:
                embedding_matrix = np.append(embedding_matrix, embedding, 0)

    return c, embedding_matrix
