# Author: Kyle Lesack
# Date: August 1st, 2024

import argparse
from Bio import SeqIO # FASTA sequences processed using Biopython
import pandas as pd # pandas used to work with csvs and tabular data
from pathlib import Path # Read and write to disk 

DEFAULT_NUM_AAS = 10 # Number of preceding amino acids to extract if not specified by user
DEFAULT_NTERM_AAS_EXCLUDE = 40 # Number of amino acids to exclude from the n-terminal sequence if not specified by user

parser = argparse.ArgumentParser()
parser.add_argument("input_csv", help = "CSV file containing two columns: (1) the FASTA file headers and (2) the peptide sequence to query against the FASTA file")
parser.add_argument("input_fasta", help = "FASTA reference file")
parser.add_argument("output_prefix", help = "Prefix for the output files")
parser.add_argument("-n", "--number_aas", type=int, help="Number of preceding amino acids to extract from FASTA file")
args = parser.parse_args()

# Function to get the amino acid sequence that precedes the query sequence
def get_preceding_aas(fasta_df, fasta_header, peptide_seq, preceding_aa_number):
	aas_extracted = False # Boolean to represent if the preceding amino acids were extracted successfully
	if fasta_header in fasta_df.keys():
		fasta_record = fasta_df[fasta_header] # Get the FASTA record from the FASTA header
		sequence_hits = fasta_record.seq.count_overlap(peptide_seq) # Number of times the sequence is found in the AA sequence
		if sequence_hits == 1:
			target_start_coord = fasta_record.seq.find(peptide_seq) # Search for the query sequence
			preceding_start_coord = target_start_coord - preceding_aa_number # Start coordinate for the preceding amino acid sequence
			if preceding_start_coord >= 0:
				results = fasta_record.seq[preceding_start_coord:target_start_coord] # Get the preceding amino acid sequence
				aas_extracted = True
			elif preceding_start_coord < 0:
				results = "Target sequence is not preceded by the required amount of amino acids. Target sequence start coordinate = " + str(target_start_coord)

		elif sequence_hits > 1:
			results = "Query sequence matched > 1 target subsequences"

		elif sequence_hits == 0:
			results = "Query sequence did not match any target subsequences"
	else:
		results = "FASTA header not found"

	return (aas_extracted, results)

# Function to replace the FASTA headers to match those in the csv files
def fix_fasta_headers(fasta_df): 
	fasta_headers = fasta_df.keys()
	new_header_dict = {}

	for header in fasta_headers: 
		new_header = header.split("|")[1] # New header is the 2nd item in the split
		new_header_dict[header] = new_header # Add new header to dictionary

	for k, v in list(fasta_df.items()): # Replace the old headers with the new ones
		fasta_df[new_header_dict.get(k, k)] = fasta_df.pop(k)
	
input_csv_file = Path(args.input_csv) # Read the CSV file containing the FASTA header and search sequence
input_fasta_file = Path(args.input_fasta) # FASTA file with target sequences

if not input_csv_file.is_file():
	print("Error: Could not find the csv file: " + args.input_csv)
elif not input_fasta_file.is_file():
	print("Error: Could not find the fasta file: " + args.input_fasta)
else:
	csv_df = pd.read_csv(args.input_csv, header = None)
	csv_df.columns = ["Header", "Peptide Sequence"]
	successful_aa_extraction = {} # Dictionary to store queries that successfully returned the preceding amino acids
	failed_aa_extraction = {} # Dictionary to store failed queries
	
	fasta_df = SeqIO.to_dict(SeqIO.parse(args.input_fasta, "fasta")) # Convert the input FASTA file to DataFrame
	fix_fasta_headers(fasta_df)

	if args.number_aas: # If the user specified a number of amino acids to extract, use it
		preceding_aa_number = args.number_aas
	else:
		print("Number of preceding amino acids to extract not provided. Using the default of " + str(DEFAULT_NUM_AAS))
		preceding_aa_number = DEFAULT_NUM_AAS
		
	for row in csv_df.iterrows(): # Iterate throught the CSV file
		fasta_header = row[1]['Header']
		peptide_seq = row[1]['Peptide Sequence']

		preceding_aas = get_preceding_aas(fasta_df, fasta_header, peptide_seq, preceding_aa_number) # Get the preceding amino acids
		if preceding_aas[0] == True:
			successful_aa_extraction[(fasta_header, peptide_seq)] = [preceding_aas[1]]
		elif preceding_aas[0] == False:
			failed_aa_extraction[(fasta_header, peptide_seq)] = preceding_aas[1]


csv_df.set_index(['Header'],inplace=True) # Set the index to the header column to facilitate joining

# Create a DataFrame with the successful queries and merge it with the original CSV file, keeping only the shared headers
successful_aa_extraction_df = pd.DataFrame.from_dict(successful_aa_extraction, orient = 'index', dtype = str)
successful_aa_extraction_df = successful_aa_extraction_df.reset_index()
successful_aa_extraction_df[['Header', 'Peptide Sequence']] = pd.DataFrame(successful_aa_extraction_df['index'].tolist(), index=successful_aa_extraction_df.index)
successful_aa_extraction_df = successful_aa_extraction_df.rename(columns={0: 'Preceding AAs'})
successful_aa_extraction_df = successful_aa_extraction_df[['Header', 'Peptide Sequence', 'Preceding AAs']]
successful_aa_extraction_df = successful_aa_extraction_df.sort_values(by=['Header', 'Peptide Sequence'])

if not successful_aa_extraction_df.empty: # Make sure that DataFrame is populated before writing results to disk
	outfile_successful = args.output_prefix + "_hits.csv" # File to store results of successful queries
	print("Writing successful queries to: " + outfile_successful)
	successful_aa_extraction_df.to_csv(outfile_successful, index = False)

# Create a DataFrame with the failed queries and merge it with the original CSV file, keeping only the shared headers
failed_aa_extraction_df = pd.DataFrame.from_dict(failed_aa_extraction, orient = 'index', dtype = str)
failed_aa_extraction_df = failed_aa_extraction_df.reset_index()
failed_aa_extraction_df[['Header', 'Peptide Sequence']] = pd.DataFrame(failed_aa_extraction_df['index'].tolist(), index=failed_aa_extraction_df.index)
failed_aa_extraction_df = failed_aa_extraction_df.rename(columns={0: 'Warnings'})
failed_aa_extraction_df = failed_aa_extraction_df[['Header', 'Peptide Sequence', 'Warnings']]
failed_aa_extraction_df = failed_aa_extraction_df.sort_values(by=['Header', 'Peptide Sequence'])

if not failed_aa_extraction_df.empty: # Make sure that DataFrame is populated before writing results to disk
	outfile_failed = args.output_prefix + "_failed.csv" # File to store results of failed queries
	print("Writing failed queries to: " + outfile_failed)
	failed_aa_extraction_df.to_csv(outfile_failed, index = False)