#! /usr/bin/env python3

import sys
import matplotlib.pyplot as plt

from Bio import Phylo
from Bio import AlignIO
from Bio.Phylo.TreeConstruction import DistanceTreeConstructor
from Bio.Phylo.TreeConstruction import DistanceCalculator

def main():
    fname = 'simple-animals.fa'
    if len(sys.argv) > 1:
        fname = sys.argv[1]
    file_format = 'fasta'
    if fname.endswith('.gb') or fname.endswith('.gbk'):
        file_format = 'genbank'
    print(fname, file_format)
    tree_alignment = AlignIO.read(fname, file_format)
    calculator = DistanceCalculator('identity')
    distance_matrix = calculator.get_distance(tree_alignment)
    constructor = DistanceTreeConstructor(calculator, 'upgma')
    tree = constructor.upgma(distance_matrix)
    Phylo.draw(tree)
    if fname == 'simple-animals.fa':
        tree.root_with_outgroup('Lamprey') # maybe not?
    Phylo.draw(tree, do_show=False)
    plt.savefig('inferred_tree.png')
    Phylo.draw_ascii(tree)

main()