import argparse
import os
import re


#Read a FASTA file generated by ProteinMPN
def read_fasta(fasta_file):
  # Open the FASTA file for reading
  with open(fasta_file, "r") as handle:
    # Initialize variables to store entries
    entries = []
    current_entry = None
    for line in handle:
      line = line.rstrip()
      # Check for header line
      if line.startswith(">"):
        # Create a new entry if needed
        if current_entry:
          entries.append(current_entry)
        current_entry = {"header": line[1:], "sequence": ""}
      else:
        # Append sequence lines to the current entry
        current_entry["sequence"] += line
  
  return entries


# Define a function to extract score from the sequence ID
def get_score(header):
  score_str = re.search(r" score=(\d+\.\d+)", header).group(1)
  return float(score_str)

# Extract each of the num_records best sequences (default 10) and write to the output file
def find_top_sequences(output_file, entries, num_records=10):
  top_sequences = []
  with open(output_file, "w") as handle:
    # Extract the first num_records sequences (lowest scores) and also write to the output file
    for entry in entries[:num_records]:
      top_sequences.append(entry['sequence'])
      handle.write(f">{entry['header']}\n{entry['sequence']}\n")
  return top_sequences
  
#Calculate consensus sequence of best num_records sequences
def calculate_consensus(top_sequences):
  consensus = ""
  for i in range(min(len(seq) for seq in top_sequences)):  # Iterate to the shortest length - although all sequences should have the same length
    aa_at_current_position = [seq[i] for seq in top_sequences]
    # Find most common amino acid
    most_common = max(set(aa_at_current_position), key=aa_at_current_position.count)  
    consensus += most_common

  return consensus
  
  
  
def main():
  """Processes all FASTA files and combines consensus sequences into a single file."""
  # Parse arguments
  parser = argparse.ArgumentParser(description="Calculate consensus sequences from FASTA files")
  parser.add_argument("input_file", help="Path to the FASTA file")
  parser.add_argument("output_file", help="Path to the output file for combined consensus sequences")
  parser.add_argument("num_for_consensus", help="Path to the output file for combined consensus sequences", nargs='?', default=10,type=int)
  args = parser.parse_args()

  #read fasta file
  entries = read_fasta(args.input_file)
  
  # Check if any entries were parsed
  if not entries:
    print(f"No FASTA entries found in {fasta_file}!")
    return
  
  # Sort entries by score in ascending order
  entries.sort(key=lambda entry: get_score(entry["header"]))
  # Find the best num_for_consensus sequences and extract to output file
  top_sequences=find_top_sequences(args.output_file, entries, args.num_for_consensus)
  #Calculate the consensus of the best sequences
  consensus_sequence=calculate_consensus(top_sequences)
  print (f"Consensus of best {args.num_for_consensus} sequences:\n{consensus_sequence}")
  

if __name__ == "__main__":
  main()

