#!/bin/python

usage="""
python phylotree.py <nwk file> [output format]
    <nwk file>      File path of nwk tree
    [output format] Output format(optional)
                    avaliable value:
                        png - png format(default)
                        svg - svg format
                        pdf - pdf format
                        all - all three format
""".strip()

from Bio        import Phylo
from matplotlib import pyplot as plt
import sys

if(len(sys.argv)<2):
    print(usage)
    sys.exit(0)
if(sys.argv[1].lower() in ["--help","-h","/?","/help"]):
    print(usage)
    sys.exit(0)

fpath = sys.argv[1]
if(len(sys.argv)==3):
    formt = sys.argv[2]
else:
    formt = "png"
name = ".".join((fpath.split('/')[-1].split('\\')[-1].split("."))[0:-1])

# read tree
tree  = Phylo.read(fpath,"newick")

# set fig size
n = len(tree.get_terminals()) # node number
figsize_x = 0.07*n+15.43
figsize_y = 0.20*n+2.19
plt.rcParams["figure.figsize"]=[figsize_x,figsize_y]  # set fig size
plt.rcParams["figure.dpi"    ]=200                    # set fig dpi
Phylo.draw(tree,do_show=False)

# set fig title
plt.title(name)

# save fig
svg=False
pdf=False
png=False
if(formt.strip().lower()=="svg"):
    svg=True
elif(formt.strip().lower()=="pdf"):
    pdf=True
elif(formt.strip().lower()=="png"):
    png=True
elif(formt.strip().lower()=="all"):
    svg=True
    pdf=True
    png=True
else:
    png=True # default: save as png

if(svg):
    plt.savefig("{}.svg".format(name))
if(pdf):
    plt.savefig("{}.pdf".format(name))
if(png):
    plt.savefig("{}.png".format(name))