import os
import csv
from Bio import SeqIO

input_folder = os.getcwd()
output_csv = os.path.join(input_folder, 'endolysin_CDS_count.csv')

with open(output_csv, mode='w', newline='') as csvfile:
    fieldnames = ['File', 'Endolysin_CDS_Exists', 'Endolysin_CDS_Count']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()
    
    for filename in os.listdir(input_folder):
        if filename.endswith('.gbk'):
            file_path = os.path.join(input_folder, filename)
            endolysin_count = 0
            endolysin_found = False
            
            for record in SeqIO.parse(file_path, 'genbank'):
                for feature in record.features:
                    if feature.type == 'CDS' and 'product' in feature.qualifiers:
                        if 'endolysin' in feature.qualifiers['product'][0].lower():
                            endolysin_count += 1
                            endolysin_found = True
            
            writer.writerow({
                'File': filename,
                'Endolysin_CDS_Exists': 1 if endolysin_found else 0,
                'Endolysin_CDS_Count': endolysin_count
            })

print(f"Results have been saved to {output_csv}")