#!/usr/bin/env python3
"""
Tree operations helper script for common ETE toolkit tasks.

Provides command-line interface for basic tree operations like:
- Format conversion
- Rooting (outgroup, midpoint)
- Pruning
- Basic statistics
- ASCII visualization
"""

import argparse
import sys
from pathlib import Path

try:
    from ete3 import Tree
except ImportError:
    print("Error: ete3 not installed. Install with: pip install ete3")
    sys.exit(1)


def load_tree(tree_file, format_num=0):
    """Load tree from file."""
    try:
        return Tree(str(tree_file), format=format_num)
    except Exception as e:
        print(f"Error loading tree: {e}")
        sys.exit(1)


def convert_format(tree_file, output, in_format=0, out_format=1):
    """Convert tree between Newick formats."""
    tree = load_tree(tree_file, in_format)
    tree.write(outfile=str(output), format=out_format)
    print(f"Converted {tree_file} (format {in_format}) → {output} (format {out_format})")


def reroot_tree(tree_file, output, outgroup=None, midpoint=False, format_num=0):
    """Reroot tree by outgroup or midpoint."""
    tree = load_tree(tree_file, format_num)

    if midpoint:
        midpoint_node = tree.get_midpoint_outgroup()
        tree.set_outgroup(midpoint_node)
        print(f"Rerooted tree using midpoint method")
    elif outgroup:
        try:
            outgroup_node = tree & outgroup
            tree.set_outgroup(outgroup_node)
            print(f"Rerooted tree using outgroup: {outgroup}")
        except Exception as e:
            print(f"Error: Could not find outgroup '{outgroup}': {e}")
            sys.exit(1)
    else:
        print("Error: Must specify either --outgroup or --midpoint")
        sys.exit(1)

    tree.write(outfile=str(output), format=format_num)
    print(f"Saved rerooted tree to: {output}")


def prune_tree(tree_file, output, keep_taxa, preserve_length=True, format_num=0):
    """Prune tree to keep only specified taxa."""
    tree = load_tree(tree_file, format_num)

    # Read taxa list
    taxa_file = Path(keep_taxa)
    if taxa_file.exists():
        with open(taxa_file) as f:
            taxa = [line.strip() for line in f if line.strip()]
    else:
        taxa = [t.strip() for t in keep_taxa.split(",")]

    print(f"Pruning tree to {len(taxa)} taxa")

    try:
        tree.prune(taxa, preserve_branch_length=preserve_length)
        tree.write(outfile=str(output), format=format_num)
        print(f"Pruned tree saved to: {output}")
        print(f"Retained {len(tree)} leaves")
    except Exception as e:
        print(f"Error pruning tree: {e}")
        sys.exit(1)


def tree_stats(tree_file, format_num=0):
    """Display tree statistics."""
    tree = load_tree(tree_file, format_num)

    print(f"\n=== Tree Statistics ===")
    print(f"File: {tree_file}")
    print(f"Number of leaves: {len(tree)}")
    print(f"Total nodes: {len(list(tree.traverse()))}")

    farthest_leaf, distance = tree.get_farthest_leaf()
    print(f"Tree depth: {distance:.4f}")
    print(f"Farthest leaf: {farthest_leaf.name}")

    # Branch length statistics
    branch_lengths = [node.dist for node in tree.traverse() if not node.is_root()]
    if branch_lengths:
        print(f"\nBranch length statistics:")
        print(f"  Mean: {sum(branch_lengths)/len(branch_lengths):.4f}")
        print(f"  Min: {min(branch_lengths):.4f}")
        print(f"  Max: {max(branch_lengths):.4f}")

    # Support values
    supports = [node.support for node in tree.traverse() if not node.is_leaf() and hasattr(node, 'support')]
    if supports:
        print(f"\nSupport value statistics:")
        print(f"  Mean: {sum(supports)/len(supports):.2f}")
        print(f"  Min: {min(supports):.2f}")
        print(f"  Max: {max(supports):.2f}")

    print()


def show_ascii(tree_file, format_num=0, show_internal=True):
    """Display tree as ASCII art."""
    tree = load_tree(tree_file, format_num)
    print(tree.get_ascii(show_internal=show_internal))


def list_leaves(tree_file, format_num=0):
    """List all leaf names."""
    tree = load_tree(tree_file, format_num)
    for leaf in tree:
        print(leaf.name)


def main():
    parser = argparse.ArgumentParser(
        description="ETE toolkit tree operations helper",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Convert format
  %(prog)s convert input.nw output.nw --in-format 0 --out-format 1

  # Midpoint root
  %(prog)s reroot input.nw output.nw --midpoint

  # Reroot with outgroup
  %(prog)s reroot input.nw output.nw --outgroup "Outgroup_species"

  # Prune tree
  %(prog)s prune input.nw output.nw --keep-taxa "speciesA,speciesB,speciesC"

  # Show statistics
  %(prog)s stats input.nw

  # Display as ASCII
  %(prog)s ascii input.nw

  # List all leaves
  %(prog)s leaves input.nw
        """
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # Convert command
    convert_parser = subparsers.add_parser("convert", help="Convert tree format")
    convert_parser.add_argument("input", help="Input tree file")
    convert_parser.add_argument("output", help="Output tree file")
    convert_parser.add_argument("--in-format", type=int, default=0, help="Input format (default: 0)")
    convert_parser.add_argument("--out-format", type=int, default=1, help="Output format (default: 1)")

    # Reroot command
    reroot_parser = subparsers.add_parser("reroot", help="Reroot tree")
    reroot_parser.add_argument("input", help="Input tree file")
    reroot_parser.add_argument("output", help="Output tree file")
    reroot_parser.add_argument("--outgroup", help="Outgroup taxon name")
    reroot_parser.add_argument("--midpoint", action="store_true", help="Use midpoint rooting")
    reroot_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")

    # Prune command
    prune_parser = subparsers.add_parser("prune", help="Prune tree to specified taxa")
    prune_parser.add_argument("input", help="Input tree file")
    prune_parser.add_argument("output", help="Output tree file")
    prune_parser.add_argument("--keep-taxa", required=True,
                              help="Taxa to keep (comma-separated or file path)")
    prune_parser.add_argument("--no-preserve-length", action="store_true",
                              help="Don't preserve branch lengths")
    prune_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")

    # Stats command
    stats_parser = subparsers.add_parser("stats", help="Display tree statistics")
    stats_parser.add_argument("input", help="Input tree file")
    stats_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")

    # ASCII command
    ascii_parser = subparsers.add_parser("ascii", help="Display tree as ASCII art")
    ascii_parser.add_argument("input", help="Input tree file")
    ascii_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")
    ascii_parser.add_argument("--no-internal", action="store_true",
                              help="Don't show internal node names")

    # Leaves command
    leaves_parser = subparsers.add_parser("leaves", help="List all leaf names")
    leaves_parser.add_argument("input", help="Input tree file")
    leaves_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    # Execute command
    if args.command == "convert":
        convert_format(args.input, args.output, args.in_format, args.out_format)
    elif args.command == "reroot":
        reroot_tree(args.input, args.output, args.outgroup, args.midpoint, args.format)
    elif args.command == "prune":
        prune_tree(args.input, args.output, args.keep_taxa,
                   not args.no_preserve_length, args.format)
    elif args.command == "stats":
        tree_stats(args.input, args.format)
    elif args.command == "ascii":
        show_ascii(args.input, args.format, not args.no_internal)
    elif args.command == "leaves":
        list_leaves(args.input, args.format)


if __name__ == "__main__":
    main()