# Path: augur/align.py
# class AlignmentError(Exception):
# def register_arguments(parser):
# def prepare(sequences, existing_aln_fname, output, ref_name, ref_seq_fname):
# def run(args):
# def postprocess(output_file, ref_name, keep_reference, fill_gaps):
# def read_sequences(*fnames):
# def check_arguments(args):
# def read_alignment(fname):
# def ensure_reference_strain_present(ref_name, existing_alignment, seqs):
# def read_reference(ref_fname):
# def generate_alignment_cmd(method, nthreads, existing_aln_fname, seqs_to_align_fname, aln_fname, log_fname):
# def remove_reference_sequence(seqs, reference_name):
# def strip_non_reference(aln, reference, insertion_csv=None):
# def analyse_insertions(aln, ungapped, insertion_csv):
# def prettify_alignment(aln):
# def make_gaps_ambiguous(aln):
# def check_duplicates(*values):
#     def add(name):
# def write_seqs(seqs, fname):
# def prune_seqs_matching_alignment(seqs, aln):


import argparse
import functools
import os
import pytest
import pathlib
from shlex import quote
from Bio import SeqIO
from Bio.Align import MultipleSeqAlignment
from Bio.Seq import Seq
from Bio.Seq import MutableSeq
from Bio.SeqRecord import SeqRecord
from augur import align
        from augur.align import prepare as prpr # Prevent recursion depth exceeded. Yes, this is terrible.
        from augur.align import prepare as prpr # Prevent recursion depth exceeded. Yes, this is terrible.
        assert ref_name == ref_seq.id, "Didn't return strain name from refrence file"
        assert os.path.isfile(aln_outfile), "Didn't write existing alignment where it said"
        assert aln_outfile != existing_aln, "Unexpectedly overwrote existing alignment"
        # Alignment file should have the reference added
        aln_output = SeqIO.to_dict(SeqIO.parse(aln_outfile, "fasta"))
        assert aln_output[ref_seq.id].seq == ref_seq.seq, "Reference sequence not added to alignment"
        for seq in existing_aln:
            assert seq in aln_output, "Some existing alignment sequences dropped unexpectedly"
            assert aln_output[seq].seq == existing_aln[seq].seq, "Some existing alignment sequences changed unexpectedly"
        # test sequences should be unchanged
        assert os.path.isfile(seqs_outfile), "Didn't write test sequences where it said"
        seq_output = SeqIO.to_dict(SeqIO.parse(seqs_outfile, "fasta"))
        for seq in test_seqs:
            assert seq in seq_output, "Some test sequences unexpectedly dropped"
            assert seq_output[seq].seq == test_seqs[seq].seq, "Some test sequences changed unexpectedly"
        assert seq_output.keys() == test_seqs.keys()
    
    def test_prepare_no_alignment_multiple_test_seqs(self, test_file, test_seqs, ref_file, ref_seq, out_file):
        """Test that we can pass multiple sequence files to prepare() and get one unified file back"""
        # bit of a kludge, but gets us the reference strain in the input files
        _, seq_outfile, _ = align.prepare([test_file, ref_file], None, out_file, ref_seq.id, None)
        seq_output = SeqIO.to_dict(SeqIO.parse(seq_outfile, "fasta"))
        assert seq_output.keys() == set(test_seqs.keys()) | {ref_seq.id}, "Did not combine the two files"
        assert seq_output[ref_seq.id].seq == ref_seq.seq, "Missing sequence from second file"
        for seq in test_seqs:
            assert seq in seq_output, "Some test sequences unexpectedly dropped"
            assert seq_output[seq].seq == test_seqs[seq].seq, "Some test sequences changed unexpectedly"
    
    def test_prepare_with_alignment_with_duplicate_sequences(self, test_file, test_seqs, existing_file, existing_aln, out_file):
        """Test that sequences matching the alignment are removed from the input sequences"""

# Path: augur/align.py
# class AlignmentError(Exception):
# def register_arguments(parser):
# def prepare(sequences, existing_aln_fname, output, ref_name, ref_seq_fname):
# def run(args):
# def postprocess(output_file, ref_name, keep_reference, fill_gaps):
# def read_sequences(*fnames):
# def check_arguments(args):
# def read_alignment(fname):
# def ensure_reference_strain_present(ref_name, existing_alignment, seqs):
# def read_reference(ref_fname):
# def generate_alignment_cmd(method, nthreads, existing_aln_fname, seqs_to_align_fname, aln_fname, log_fname):
# def remove_reference_sequence(seqs, reference_name):
# def strip_non_reference(aln, reference, insertion_csv=None):
# def analyse_insertions(aln, ungapped, insertion_csv):
# def prettify_alignment(aln):
# def make_gaps_ambiguous(aln):
# def check_duplicates(*values):
#     def add(name):
# def write_seqs(seqs, fname):
# def prune_seqs_matching_alignment(seqs, aln):


# Path: tests/test_align.py
import argparse
import functools
import os
import pytest
import pathlib
from shlex import quote
from Bio import SeqIO
from Bio.Align import MultipleSeqAlignment
from Bio.Seq import Seq
from Bio.Seq import MutableSeq
from Bio.SeqRecord import SeqRecord
from augur import align
        from augur.align import prepare as prpr # Prevent recursion depth exceeded. Yes, this is terrible.
        from augur.align import prepare as prpr # Prevent recursion depth exceeded. Yes, this is terrible.

        assert ref_name == ref_seq.id, "Didn't return strain name from refrence file"
        assert os.path.isfile(aln_outfile), "Didn't write existing alignment where it said"
        assert aln_outfile != existing_aln, "Unexpectedly overwrote existing alignment"
        # Alignment file should have the reference added
        aln_output = SeqIO.to_dict(SeqIO.parse(aln_outfile, "fasta"))
        assert aln_output[ref_seq.id].seq == ref_seq.seq, "Reference sequence not added to alignment"
        for seq in existing_aln:
            assert seq in aln_output, "Some existing alignment sequences dropped unexpectedly"
            assert aln_output[seq].seq == existing_aln[seq].seq, "Some existing alignment sequences changed unexpectedly"
        # test sequences should be unchanged
        assert os.path.isfile(seqs_outfile), "Didn't write test sequences where it said"
        seq_output = SeqIO.to_dict(SeqIO.parse(seqs_outfile, "fasta"))
        for seq in test_seqs:
            assert seq in seq_output, "Some test sequences unexpectedly dropped"
            assert seq_output[seq].seq == test_seqs[seq].seq, "Some test sequences changed unexpectedly"
        assert seq_output.keys() == test_seqs.keys()
    
    def test_prepare_no_alignment_multiple_test_seqs(self, test_file, test_seqs, ref_file, ref_seq, out_file):
        """Test that we can pass multiple sequence files to prepare() and get one unified file back"""
        # bit of a kludge, but gets us the reference strain in the input files
        _, seq_outfile, _ = align.prepare([test_file, ref_file], None, out_file, ref_seq.id, None)
        seq_output = SeqIO.to_dict(SeqIO.parse(seq_outfile, "fasta"))
        assert seq_output.keys() == set(test_seqs.keys()) | {ref_seq.id}, "Did not combine the two files"
        assert seq_output[ref_seq.id].seq == ref_seq.seq, "Missing sequence from second file"
        for seq in test_seqs:
            assert seq in seq_output, "Some test sequences unexpectedly dropped"
            assert seq_output[seq].seq == test_seqs[seq].seq, "Some test sequences changed unexpectedly"
    
    def test_prepare_with_alignment_with_duplicate_sequences(self, test_file, test_seqs, existing_file, existing_aln, out_file):
        """Test that sequences matching the alignment are removed from the input sequences"""

        _, seq_outfile, _ = align.prepare([existing_file, test_file], existing_file, out_file, None, None)