import sys
import os
from Bio import SeqIO
from Bio.SeqFeature import SeqFeature

# Script to insert gene features to gbk-format file
#
# Usage: python add_gene_features_to_gbk.py annotation.gbk
#
# input: `annotation.gbk`, GenBank format file generated by DFAST 
# output: `annotation.with_gene.gbk`, GenBank format file with gene features inserted


input_file = sys.argv[1]

in_base, in_ext = os.path.splitext(input_file)
output_file = f"{in_base}.with_gene{in_ext}"

R = list(SeqIO.parse(input_file, "genbank"))
for r in R:
    features = []
    for f in r.features:
        r.description = ""
        if "locus_tag" in f.qualifiers:
            locus_tag = f.qualifiers["locus_tag"][0]
            gene = f.qualifiers.get("gene", [""])[0]
            qualifiers = {"locus_tag": [locus_tag]}
            if gene:
                qualifiers["gene"] = [gene]
            gene_feature = SeqFeature(location=f.location, type="gene", qualifiers=qualifiers)
            features.append(gene_feature)
        features.append(f)
    r.features = features

with open(output_file, "w") as f:
    SeqIO.write(R, f, "genbank")

