import os
from utils.protein_group_sql_db import ProteinGroupSqlDbInteractor

checkpoint create_pg_fasta:
    input:
        fasta='seqs/all.faa',
        pg_db='pgs/protein_groups_db.sql',
    output:
        directory('seqs/pgs')  # '/pg.{i}.faa'
    message:
        "creating FASTA files for protein groups ..."
    version:
        "0.3"
    run:
        os.mkdir(output[0])
        pg_obj = ProteinGroupSqlDbInteractor(input.pg_db)
        pg_to_acc_str_dict = pg_obj.get_pg_id_to_accs_dict()
        sep = pg_obj.get_protein_sep()
        group_ids = pg_obj.get_groups_to_align_from_pg_db()
        accs = []
        for group_id in group_ids:
            accs.append(set(pg_to_acc_str_dict[group_id].split(sep)))
        started_groups = set()
        with open(input.fasta,"r") as inhandle:
            for line in inhandle:
                if len(line.strip()) == 0:
                    continue
                acc_line = line
                acc = acc_line[1:-1]
                seq = inhandle.readline()
                for acc_set_i, group_id in zip(accs,group_ids):
                    if acc in acc_set_i:
                        if not group_id in started_groups:
                            started_groups.add(group_id)
                            mode = "w"
                        else:
                            mode = "a"
                        with open(os.path.join(output[0],f'pg.{group_id}.faa'),mode) as outhandle:
                            outhandle.write(acc_line + seq)
