From 81eac9255f6fd219159b5d47cbf2d514e6546ab0 Mon Sep 17 00:00:00 2001
From: James Hadfield <hadfield.james@gmail.com>
Date: Wed, 13 Dec 2023 16:28:04 +1300
Subject: [PATCH] WIP - export root AA sequences for VCF inputs

---
 augur/translate.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/augur/translate.py b/augur/translate.py
index 0d04a56f2..9db2396ac 100644
--- a/augur/translate.py
+++ b/augur/translate.py
@@ -268,6 +268,15 @@ def assign_aa_vcf(tree, translations):
 
                 aa_muts[c.name]["aa_muts"][fname] = tmp
 
+    ## NEEDS TESTS ADDED WITH AA VARIATION AT ROOT!!!!
+    aa_muts[root.name]['aa_sequences'] = {}
+    for gene_name, gene_data in translations.items():
+        root_seq = list(gene_data['reference'])
+        for pos,alt in gene_data['sequences'][root.name].items():
+            # pos is 0-based, <class 'numpy.int64'>
+            root_seq[pos] = alt
+        aa_muts[root.name]['aa_sequences'][gene_name] = "".join(root_seq)
+
     return aa_muts
 
 def assign_aa_fasta(tree, translations):