Visualization and Output

def draw(self, show_plot: bool = True, export_filename: str = None, show_labels: bool = False, align_labels: bool = False, label_fontsize: int = 8, start_time: float = 0, default_color: str = '#000000', xlabel: str = None, handles=None) -> None:
    """
    Draw tree visualization using matplotlib.

    Parameters:
    - show_plot (bool): Display the plot window
    - export_filename (str): Save plot to file (format determined by extension)
    - show_labels (bool): Display node labels
    - align_labels (bool): Align leaf labels vertically
    - label_fontsize (int): Font size for labels
    - start_time (float): Starting time for root
    - default_color (str): Default color for tree elements
    - xlabel (str): X-axis label
    - handles: Custom legend handles
    """

import treeswift
import matplotlib.pyplot as plt

# Basic tree drawing
tree = treeswift.read_tree_newick("((A:0.1,B:0.2):0.3,(C:0.4,D:0.5):0.6);")

# Simple tree plot
tree.draw(show_labels=True)

# Customized tree plot
tree.draw(
    show_labels=True,
    align_labels=True,
    label_fontsize=12,
    default_color='blue',
    xlabel='Branch Length',
    export_filename='my_tree.png'
)

# Multiple trees on same figure
tree1 = treeswift.read_tree_newick("((A,B),(C,D));")
tree2 = treeswift.read_tree_newick("(((A,B),C),D);")

plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)
tree1.draw(show_plot=False, show_labels=True)
plt.title("Balanced Tree")

plt.subplot(1, 2, 2)
tree2.draw(show_plot=False, show_labels=True)
plt.title("Unbalanced Tree")

plt.tight_layout()
plt.show()

def lineages_through_time(self, present_day: float = None, show_plot: bool = True, export_filename: str = None, color: str = '#000000', xmin: float = None, xmax: float = None, ymin: float = None, ymax: float = None, title: str = None, xlabel: str = None, ylabel: str = None) -> dict:
    """
    Compute and optionally plot lineages through time.

    Parameters:
    - present_day (float): Time point representing present day
    - show_plot (bool): Display the LTT plot
    - export_filename (str): Save plot to file
    - color (str): Plot line color
    - xmin, xmax (float): X-axis limits
    - ymin, ymax (float): Y-axis limits
    - title (str): Plot title
    - xlabel (str): X-axis label
    - ylabel (str): Y-axis label

    Returns:
    - dict: Dictionary mapping times to lineage counts
    """

def ltt(self, **kwargs) -> dict:
    """Alias for lineages_through_time with same parameters."""

def plot_ltt(lineages: dict, show_plot: bool = True, export_filename: str = None, color: str = '#000000', xmin: float = None, xmax: float = None, ymin: float = None, ymax: float = None, title: str = None, xlabel: str = None, ylabel: str = None) -> None:
    """
    Plot lineages through time from precomputed data.

    Parameters:
    - lineages (dict): Lineages dictionary from lineages_through_time()
    - show_plot (bool): Display the plot
    - export_filename (str): Save plot to file
    - color (str): Plot line color
    - xmin, xmax (float): X-axis limits
    - ymin, ymax (float): Y-axis limits
    - title (str): Plot title
    - xlabel (str): X-axis label
    - ylabel (str): Y-axis label
    """

import treeswift

# Basic LTT analysis
tree = treeswift.read_tree_newick("((A:0.1,B:0.2):0.3,(C:0.4,D:0.5):0.6);")

# Generate LTT plot
ltt_data = tree.lineages_through_time(
    show_plot=True,
    title="Lineages Through Time",
    xlabel="Time (My)",
    ylabel="Number of Lineages"
)

print("LTT data points:")
for time, lineages in sorted(ltt_data.items()):
    print(f"  Time {time:.3f}: {lineages} lineages")

# Customized LTT plot
tree.ltt(
    color='red',
    title='Diversification Pattern',
    xlabel='Time before present',
    ylabel='Lineage count',
    export_filename='ltt_plot.pdf'
)

# Compare multiple trees
trees = [
    treeswift.read_tree_newick("((A:0.1,B:0.1):0.8,(C:0.1,D:0.1):0.8);"),  # Early burst
    treeswift.read_tree_newick("((A:0.8,B:0.8):0.1,(C:0.8,D:0.8):0.1);"),  # Late burst
]

colors = ['blue', 'red']
labels = ['Early diversification', 'Late diversification']

plt.figure(figsize=(10, 6))
for i, tree in enumerate(trees):
    ltt_data = tree.lineages_through_time(show_plot=False)
    treeswift.plot_ltt(
        ltt_data,
        show_plot=False,
        color=colors[i]
    )

plt.legend(labels)
plt.title('Comparison of Diversification Patterns')
plt.xlabel('Time')
plt.ylabel('Number of Lineages')
plt.show()

def newick(self) -> str:
    """
    Generate Newick format string.

    Returns:
    - str: Standard Newick representation of tree
    """

def indent(self, space: int = 4) -> str:
    """
    Generate indented text representation (like nw_indent).

    Parameters:
    - space (int): Number of spaces per indentation level

    Returns:
    - str: Indented tree representation
    """

import treeswift

tree = treeswift.read_tree_newick("((A:0.1,B:0.2):0.3,(C:0.4,D:0.5):0.6);")

# Standard Newick output
newick_str = tree.newick()
print("Newick format:")
print(newick_str)

# Indented representation
indented = tree.indent(space=2)
print("\nIndented format:")
print(indented)

# More readable indentation
readable = tree.indent(space=4)
print("\nReadable format:")
print(readable)

# Compare formats for complex tree
complex_tree = treeswift.read_tree_newick("(((((A,B),C),D),E),((F,G),(H,I)));")
print("\nComplex tree - Newick:")
print(complex_tree.newick())
print("\nComplex tree - Indented:")
print(complex_tree.indent())

def write_tree_newick(self, filename: str, hide_rooted_prefix: bool = False) -> None:
    """
    Write tree to Newick format file.

    Parameters:
    - filename (str): Output filename (.gz extension for compression)
    - hide_rooted_prefix (bool): Hide rooted tree prefix in output
    """

def write_tree_nexus(self, filename: str) -> None:
    """
    Write tree to Nexus format file.

    Parameters:
    - filename (str): Output filename (.gz extension for compression)
    """

import treeswift
import os

tree = treeswift.read_tree_newick("((A:0.1,B:0.2):0.3,(C:0.4,D:0.5):0.6);")

# Write to plain text files
tree.write_tree_newick("output.nwk")
tree.write_tree_nexus("output.nexus")

# Write to compressed files
tree.write_tree_newick("output.nwk.gz")
tree.write_tree_nexus("output.nexus.gz")

# Write with options
tree.write_tree_newick("unrooted_format.nwk", hide_rooted_prefix=True)

# Verify files were created
for filename in ["output.nwk", "output.nexus", "output.nwk.gz", "output.nexus.gz"]:
    if os.path.exists(filename):
        size = os.path.getsize(filename)
        print(f"{filename}: {size} bytes")

# Read back and verify
tree_copy = treeswift.read_tree_newick("output.nwk")
print(f"Original: {tree.newick()}")
print(f"Round-trip: {tree_copy.newick()}")
print(f"Identical: {tree.newick() == tree_copy.newick()}")

# Additional visualization parameters for draw() method
def draw(self, 
         # ... basic parameters ...
         show_plot: bool = True,
         export_filename: str = None,
         # Label options
         show_labels: bool = False,
         align_labels: bool = False, 
         label_fontsize: int = 8,
         # Time and color options
         start_time: float = 0,
         default_color: str = '#000000',
         # Axis options  
         xlabel: str = None,
         # Legend options
         handles = None
         ) -> None:
    """Extended visualization with advanced customization."""

import treeswift
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Advanced visualization example
tree = treeswift.read_tree_newick("((A:0.1,B:0.2):0.3,(C:0.4,D:0.5):0.6);")

# Create custom visualization
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# Left plot: Basic visualization
plt.sca(ax1)
tree.draw(
    show_plot=False,
    show_labels=True,
    align_labels=True,
    label_fontsize=10,
    default_color='black'
)
ax1.set_title('Standard Tree Plot')

# Right plot: Time-scaled visualization  
plt.sca(ax2)
tree.draw(
    show_plot=False,
    show_labels=True,
    align_labels=True,
    label_fontsize=10,
    start_time=1.0,  # Root at time 1.0
    default_color='darkblue',
    xlabel='Time (Ma)'
)
ax2.set_title('Time-Scaled Tree')

plt.tight_layout()
plt.show()

# Export high-quality figure
tree.draw(
    show_labels=True,
    align_labels=True,
    label_fontsize=12,
    default_color='black',
    xlabel='Evolutionary Distance',
    export_filename='publication_tree.pdf'  # Vector format for publications
)

# Export raster format
tree.draw(
    show_labels=True,
    export_filename='presentation_tree.png'  # Raster format for presentations
)

import treeswift
import matplotlib.pyplot as plt

# Visualize multiple trees in grid
tree_strings = [
    "((A,B),(C,D));",           # Balanced
    "(((A,B),C),D);",           # Unbalanced  
    "(A,B,C,D);",               # Star (polytomy)
    "((A:0.1,B:0.9):0.1,(C:0.5,D:0.5):0.5);"  # Mixed branch lengths
]

titles = ['Balanced', 'Unbalanced', 'Star Tree', 'Mixed Branch Lengths']
trees = [treeswift.read_tree_newick(s) for s in tree_strings]

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
axes = axes.flatten()

for i, (tree, title) in enumerate(zip(trees, titles)):
    plt.sca(axes[i])
    tree.draw(
        show_plot=False,
        show_labels=True,
        align_labels=True,
        label_fontsize=8
    )
    axes[i].set_title(title)

plt.tight_layout()
plt.savefig('tree_comparison.png', dpi=300, bbox_inches='tight')
plt.show()

# Create LTT comparison
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# Individual LTT plots
for i, tree in enumerate(trees[:2]):  # First two trees
    plt.sca(ax1 if i == 0 else ax2)
    ltt_data = tree.lineages_through_time(
        show_plot=False,
        title=titles[i]
    )
    treeswift.plot_ltt(ltt_data, show_plot=False)

plt.tight_layout()
plt.show()