from Bio import Phylo
import matplotlib.pyplot as plt

tree = Phylo.read('tree.nwk', 'newick')

# Quick text representation
print(tree)

# ASCII art diagram
Phylo.draw_ascii(tree)

tree = Phylo.read('tree.nwk', 'newick')

# Simple plot (opens interactive window)
Phylo.draw(tree)
plt.show()

# Save to file
fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax)
plt.savefig('tree.png', dpi=300, bbox_inches='tight')
plt.close()

fig, ax = plt.subplots(figsize=(12, 10))
Phylo.draw(tree, axes=ax, do_show=False,
           branch_labels=lambda c: f'{c.branch_length:.2f}' if c.branch_length else '',
           label_func=lambda c: c.name if c.is_terminal() else '')

ax.set_title('Phylogenetic Tree')
plt.savefig('custom_tree.png', dpi=300, bbox_inches='tight')
plt.close()

# Custom label function
def custom_labels(clade):
    if clade.is_terminal():
        return clade.name
    elif clade.confidence:
        return f'{clade.confidence:.0f}'
    return ''

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax, label_func=custom_labels)
plt.savefig('labeled_tree.png', dpi=300)
plt.close()

# Show branch lengths
def branch_length_labels(clade):
    if clade.branch_length:
        return f'{clade.branch_length:.3f}'
    return ''

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax, branch_labels=branch_length_labels)
plt.savefig('with_lengths.png', dpi=300)
plt.close()

# Show bootstrap values (stored in clade.confidence or clade.name for internal nodes)
def bootstrap_labels(clade):
    if not clade.is_terminal() and clade.confidence:
        return f'{clade.confidence:.0f}'
    return ''

Phylo.draw(tree, axes=ax, branch_labels=bootstrap_labels)

# Color specific clades before drawing
tree = Phylo.read('tree.nwk', 'newick')

# Set colors for specific clades (PhyloXML trees support this natively)
for clade in tree.find_clades():
    if clade.name and 'Human' in clade.name:
        clade.color = 'red'
    elif clade.name and 'Mouse' in clade.name:
        clade.color = 'blue'

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax)
plt.savefig('colored_tree.png', dpi=300)
plt.close()

from Bio.Phylo.PhyloXML import BranchColor

# Convert to PhyloXML for color support
phyloxml_tree = tree.as_phyloxml()

# Color a clade and its descendants
target = phyloxml_tree.find_any(name='Human')
if target:
    target.color = BranchColor.from_name('red')

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(phyloxml_tree, axes=ax)
plt.savefig('highlighted.png', dpi=300)
plt.close()

tree = Phylo.read('tree.nwk', 'newick')
tree.ladderize()

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax, do_show=False)

# PNG (raster, good for presentations)
plt.savefig('tree.png', dpi=300, bbox_inches='tight')

# PDF (vector, good for publications)
plt.savefig('tree.pdf', bbox_inches='tight')

# SVG (vector, good for web)
plt.savefig('tree.svg', bbox_inches='tight')

plt.close()

# Adjust figure size based on tree size
n_taxa = len(tree.get_terminals())
height = max(8, n_taxa * 0.3)  # Scale with number of taxa

fig, ax = plt.subplots(figsize=(10, height))
Phylo.draw(tree, axes=ax, do_show=False)
plt.tight_layout()
plt.savefig('scaled_tree.png', dpi=300)
plt.close()

tree = Phylo.read('tree.nwk', 'newick')

# Ladderize for cleaner appearance
tree.ladderize(reverse=True)

# Set missing branch lengths to small value
for clade in tree.find_clades():
    if clade.branch_length is None:
        clade.branch_length = 0.001

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax)
plt.savefig('clean_tree.png', dpi=300)
plt.close()

tree1 = Phylo.read('tree1.nwk', 'newick')
tree2 = Phylo.read('tree2.nwk', 'newick')

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))

Phylo.draw(tree1, axes=ax1, do_show=False)
ax1.set_title('Tree 1')

Phylo.draw(tree2, axes=ax2, do_show=False)
ax2.set_title('Tree 2')

plt.tight_layout()
plt.savefig('comparison.png', dpi=300)
plt.close()

fig, ax = plt.subplots(figsize=(10, 8))
Phylo.draw(tree, axes=ax, do_show=False)

ax.axis('off')  # Remove axis
ax.set_frame_on(False)  # Remove frame

plt.savefig('clean_tree.png', dpi=300, bbox_inches='tight', transparent=True)
plt.close()