from Bio import SeqIO
import pandas as pd

def genbank_url(accession):
    return f"https://www.ncbi.nlm.nih.gov/nuccore/{accession}"

# Open the GenBank file
genbank_file = "FILE-NAME.gb"  # Replace with your file path

# Create lists to store the extracted data
accession_list = []
sequence_list = []
nucleotide_size_list = []  # Updated list for nucleotide sequence size
protein_size_list = []  # New list for protein sequence size
country_list = []
host_list = []
organism_list = []  # New list for organism
protein_id_list = []
translation_list = []
pubmed_list = []

# Iterate through the GenBank file
for record in SeqIO.parse(genbank_file, "genbank"):
    # Accession code
    accession = record.id if record.id else "Missing"
    accession_list.append(accession)

    # Sequence
    sequence = str(record.seq) if record.seq else "Missing"
    sequence_list.append(sequence)

    # Nucleotide sequence size (counting the number of bases)
    nucleotide_size = sum(c.isalpha() for c in sequence)
    nucleotide_size_list.append(nucleotide_size)

    # Extracting country, host, protein ID, translation, and PUBMED if available
    country = "Missing"
    host = "Missing"
    organism = record.annotations.get('organism', 'Missing')
    protein_id = "Missing"
    translation = "Missing"
    protein_size = "Missing"
    pubmed = "Missing"
    for feature in record.features:
        if 'country' in feature.qualifiers:
            country = feature.qualifiers['country'][0]
        if 'host' in feature.qualifiers:
            host = feature.qualifiers['host'][0]
        if feature.type == 'CDS':
            if 'protein_id' in feature.qualifiers:
                protein_id = feature.qualifiers['protein_id'][0]
            if 'translation' in feature.qualifiers:
                translation = feature.qualifiers['translation'][0]
                protein_size = len(translation)  # Calculate protein size
    # Check for PUBMED references
    for reference in record.annotations['references']:
        if reference.pubmed_id:
            pubmed = reference.pubmed_id
    
    country_list.append(country)
    host_list.append(host)
    organism_list.append(organism)
    protein_id_list.append(protein_id)
    translation_list.append(translation)
    protein_size_list.append(protein_size)
    pubmed_list.append(pubmed)

# Create a pandas DataFrame
data = {
    'Accession Code': accession_list,
    'Sequence': sequence_list,
    'Nucleotide Size': nucleotide_size_list,
    'Protein Size': protein_size_list,
    'Country': country_list,
    'Host': host_list,
    'Organism': organism_list,
    'Protein ID': protein_id_list,
    'Translation': translation_list,
    'PUBMED': pubmed_list
}

df = pd.DataFrame(data)

# Add hyperlinks for Accession Code, Protein ID, and PUBMED
df['Accession Code'] = df['Accession Code'].apply(lambda x: f'=HYPERLINK("{genbank_url(x)}", "{x}")' if x != "Missing" else "Missing")
df['Protein ID'] = df['Protein ID'].apply(lambda x: f'=HYPERLINK("{genbank_url(x)}", "{x}")' if x != "Missing" else "Missing")
df['PUBMED'] = df['PUBMED'].apply(lambda x: f'=HYPERLINK("https://pubmed.ncbi.nlm.nih.gov/{x}/", "{x}")' if x != "Missing" else "Missing")

# Save to an Excel file
output_file = "output_data_hyperlinked.xlsx"  # Replace with your desired output file path
df.to_excel(output_file, index=False)