#!/usr/bin/env python
# File created on 15 Jul 2011
from __future__ import division

__author__ = "Jesse Zaneveld"
__copyright__ = "Copyright 2011-2013, The PICRUSt Project"
__credits__ = ["Jesse Zaneveld","Morgan Langille"]
__license__ = "GPL"
__version__ = "1.0.0-dev"
__maintainer__ = "Jesse Zaneveld"
__email__ = "zaneveld@gmail.com"
__status__ = "Development"

from os.path import splitext
from string import maketrans
from sys import getrecursionlimit,setrecursionlimit
import re
from cogent.parse.tree import DndParser
from cogent.util.option_parsing import parse_command_line_parameters,\
    make_option

from picrust.parse import parse_trait_table,yield_trait_table_fields
from util import PicrustNode

def reformat_tree_and_trait_table(tree,trait_table_lines,trait_to_tree_mapping,\
    input_trait_table_delimiter="\t", output_trait_table_delimiter="\t",\
    filter_table_by_tree_tips=True, convert_trait_floats_to_ints=False,\
    filter_tree_by_table_entries=True,convert_to_bifurcating=False,\
    add_branch_length_to_root=False, name_unnamed_nodes=True,\
    remove_whitespace_from_labels = True,replace_ambiguous_states=True,\
    replace_problematic_label_characters = True,min_branch_length=0.0001,\
    verbose=True):
    """Return a full reformatted tree,pruned reformatted tree  and set of trait table lines 

    tree - a PyCogent PhyloNode tree object
    
    trait_table_lines -- the lines of a trait table, where 
      the rows are organisms and the columns are traits (e.g. gene counts).
    
    trait_id_to_tree_mapping -- a dict keyed by trait table ids, with
      values of tree ids.   If provided, trait table ids will be mapped to 
      tree ids

    filter_table_by_tree_tips -- if True, remove trait table rows that don't map to ids on the 
    tree 

    convert_trait_floats_to_ints -- if True, convert floating point values in trait table cells to integers.

    filter_tree_by_table_entries -- if True, save only the subtree that encompasses organisms in the trait table.
    (equivalent to removing all tips in the tree that don't map to the trait table)

    convert_to_bifurcating -- if True, ensure that the tree is fully bifurcating by resolving polytomies with very short
    branches.

    add_branch_length_to_root -- if True, ensure that the root node has a minimum branch length

    name_unnamed_nodes -- if True, name unnamed nodes in the tree.   (Useful for ensuring internal nodes can be 
    consistently identified in both the reference and pruned trees)

    remove_whitespace_from_labels -- if True, replace whitespace in organism labels with underscores
   
    replace_ambiguous_states -- if True, replace various strings representing ambiguous character states,
    as well as '-1' or -1 (used by IMG to represent a lack of data) with 0 values.

    replace_problematic_table_chars -- if True, replace ':' and ';' in the results with '_', and remove double quotes.
    (AncSR methods like ace can't handle these characters in organism labels)

    min_branch_length -- set the minimum branch length for all edges in the tree.   
    
    This function combines the various reformatting functions in the 
    library into a catch-all reformatter.  
    
    TODO: This function is monolithic, so despite the individual
    parts being tested seperately, it probably needs to be broken
    down into several modular parts.  This would need to be done
    with care however, as the order of steps matters quite a bit.
    

    """
    
    

    input_tree = tree
    
    #Parse lines to fields once
    
    
    if trait_table_lines:
        if verbose:
            print "Parsing trait table...."
      
  header_line,trait_table_fields =\
          parse_trait_table(trait_table_lines,delimiter = input_trait_table_delimiter)
    else:
        if verbose:
            print "Found no trait table lines. Setting data and header to empty"
        trait_table_fields = []
        header_line = ''
    
    # Tree reformatting
    if convert_to_bifurcating:
        if verbose:
            print "Converting tree to bifurcating...."
       
        #maximum recursion depth on large trees
   
     #Try working around this issue with a large
        #recursion depth limit
        old_recursion_limit = getrecursionlimit()
        setrecursionlimit(50000)
        input_tree = input_tree.bifurcating() # Required by most ancSR programs
        setrecursionlimit(old_recursion_limit)

        #input_tree = ensure_root_is_bifurcating(input_tree)

        # The below nutty-looking re-filtering step is necessary
        # When ensuring the root is bifurcating, internal nodes can 
        #get moved to the tips so without additional filtering we 
        #get unannotated tip nodes
        
        #if filter_tree_by_table_entries:
        #    input_tree = filter_tree_tips_by_presence_in_table(input_tree,\
        #      trait_table_fields,delimiter=input_trait_table_delimiter)


       
    #Name unnamed nodes
    if name_unnamed_nodes:
        if verbose:
            print "Naming unnamed nodes in the reference tree...."
        input_tree=make_internal_nodes_unique(input_tree)
        #input_tree.nameUnnamedNodes()
        check_node_labels(input_tree,verbose=verbose)
        #Paranoid check for missing names:
        #if verbose:
        #    print "Checking that all nodes were named..."
        #for i,n in enumerate(input_tree.preorder()):
        #    if n.Name is None:
        #        raise ValueError('Node #%s (in tree.preorder()) was not named!'%str(i))
    
        
    #map trait table ids to tree ids
    if trait_to_tree_mapping:
        #if verbose:
        #    print "Validating that trait --> tree mappings match tree ids..."
        #    good,bad = validate_trait_table_to_tree_mappings(input_tree,\
        #      trait_to_tree_mapping.values(), verbose = True)
        #    print "Found %i valid ids." %(len(good))
        #    print "Found %i invalid ids." %(len(bad))
        #    #if bad:
        #    #    raise RuntimeError("The following putative tree ids in mapping file aren't actually in the input tree: %s" % bad)
    
    
        if verbose:
            print "Remapping trait table ids to match tree ids...."
        
        trait_table_fields =\
          remap_trait_table_organisms(trait_table_fields,trait_to_tree_mapping,\
          verbose = verbose)
    
    label_conversion_fns =\
      set_label_conversion_fns(remove_whitespace_from_labels=remove_whitespace_from_labels,\
        replace_problematic_label_characters=replace_problematic_label_characters)
    
    value_conversion_fns = set_value_conversion_fns(replace_ambiguous_states=replace_ambiguous_states,\
      convert_trait_floats_to_ints=convert_trait_floats_to_ints) 
   

    #Apply both label and value converters to the trait table
    trait_table_fields = convert_trait_table_entries(\
      trait_table_fields,\
      value_conversion_fns = value_conversion_fns,\
      label_conversion_fns = label_conversion_fns)
   
 
    #We now need to apply any formatting functions to the tree nodes as well, to ensure
    #that names are consistent between the two.
    
    if label_conversion_fns:
        input_tree = fix_tree_labels(input_tree, label_conversion_fns)
          
    #Then filter the trait table to include only tree tips
    if filter_table_by_tree_tips:
        if verbose:
            print "Filtering trait table ids to include only those that match tree ids...."
        trait_table_fields = filter_table_by_presence_in_tree(input_tree,\
          trait_table_fields,delimiter=input_trait_table_delimiter)
        
        #if verbose:
        #    print "Verifying that new trait table ids match tree:"
        #    print "# of trait_table_lines: %i" %len(trait_table_lines)
        #    all_tip_ids = [tip.Name for tip in input_tree.iterTips()]
        #    print "example tree tip ids:",all_tip_ids[0:10]
    if filter_tree_by_table_entries:
        if verbose:
            print "filtering tree tips to match entries in trait table...."
        input_tree = filter_tree_tips_by_presence_in_table(input_tree,\
          trait_table_fields,delimiter=input_trait_table_de