from Bio import SeqIO
import re
import os

def contains_consecutive_wfq(seq):
    for i in range(len(seq) - 10):  # Loop through the sequence up to the 10th-to-last position
        if re.search(r'WF.', seq[i:i+3]):
            if i >= 48 and seq[i - 36] == "Q" and seq[i - 28] == "F":
                return i
            elif i >= 48 and seq[i - 36] == "Q" and seq[i - 22] == "P":
                return i
            elif i >= 48 and seq[i - 28] == "F" and seq[i - 22] == "P":
                return i
        if re.search(r'.FQ', seq[i:i+3]):
            if i >= 48 and seq[i - 36] == "Q" and seq[i - 28] == "F":
                return i
            elif i >= 48 and seq[i - 36] == "Q" and seq[i - 22] == "P":
                return i
            elif i >= 48 and seq[i - 28] == "F" and seq[i - 22] == "P":
                return i
        if re.search(r'W.Q', seq[i:i+3]):  # Corrected: Search for the "W.Q" pattern only in the current 3-character substring
            if i >= 48 and seq[i - 36] == "Q" and seq[i - 28] == "F":
                return i
            elif i >= 48 and seq[i - 36] == "Q" and seq[i - 22] == "P":
                return i
            elif i >= 48 and seq[i - 28] == "F" and seq[i - 22] == "P":
                return i
    return 0


def write_to_file(records, output_file):
    with open(output_file, "w") as f:
        SeqIO.write(records, f, "fasta")

def filter_sequences(input_file, output_file):
    records_to_write = []
    with open(input_file, "r") as f:
        for record in SeqIO.parse(f, "fasta"):
            sequence = str(record.seq)
            i = contains_consecutive_wfq(sequence)
            if i and contains_consecutive_wfq(sequence[i+5:i+125]):
                records_to_write.append(record)
    write_to_file(records_to_write, output_file)


# Process all input files with the suffix "_genes" and output files with the suffix "_genes_DUX"
for filename in os.listdir("."):
    if filename.endswith("_genes"):
        input_file = filename
        output_file = filename.replace("_genes", "_genes_DUX")
        filter_sequences(input_file, output_file)