from __future__ import division, print_function
from Bio import Phylo
from StringIO import StringIO
import numpy as np
import sys, csv, logging, json, argparse, imp, math
from collections import defaultdict
sys.path.append('')
from base.colorLogging import ColorizingStreamHandler


version = 0.1

def get_command_line_args():
    parser = argparse.ArgumentParser(description ="Convert tree files to Auspice ready JSONs. Version {}".format(version))

    newick = parser.add_argument_group('newick')
    newick.add_argument('--newick', type=str, help="Path to newick file")

    beast = parser.add_argument_group('beast')
    beast.add_argument('--nexus', type=str, help="Path to nexus file")
    beast.add_argument('--most_recent_tip', type=float, help="Date of the most recent tip (in decimal format)")
    beast.add_argument('--discrete_traits', type=str, nargs='+', action='extend', default=[], help="Discrete traits to extract from the BEAST annotations")
    beast.add_argument('--continuous_traits', type=str, nargs='+', action='extend', default=[], help="Continuous traits to extract from the BEAST annotations")
    beast.add_argument('--make_traits_log', type=str, nargs='+', action='extend', default=[], help="Convert these (continous) traits to log space: y=-ln(x)")
    beast.add_argument("--fake_divergence", action="store_true", help="Set the divergence as time (prevents auspice crashing)")


    metadata = parser.add_argument_group('metadata')
    metadata.add_argument('--strain_csv', default=None, type=str, help="Metadata (with headers) linking strain to some metadata columns")


    general = parser.add_argument_group('general')
    general.add_argument("--maintainer", type=str, nargs=2, default=["", ""], help="Maintaner (display name, link name). Shown in auspice footer.")
    general.add_argument("--debug", action="store_const", dest="loglevel", const=logging.DEBUG, help="Enable debugging logging")
    general.add_argument('--output_prefix', '-o', required=True, type=str, help="Output prefix (i.e. \"_meta.json\" will be appended to this)")
    general.add_argument('--title', default=None, type=str, help="Title (to be displayed by auspice)")
    general.add_argument("--defaults", type=str, nargs='+', action='extend', default=[], help="Auspice defaults. Format: \"key:value\"")
    general.add_argument("--filters", type=str, nargs='+', action='extend', default=[], help="Auspice filters.")
    general.add_argument('--geo', type=str, help="CSV File w. header \"trait,value,latitude,longitude\". Turns on the map panel.")


    return parser.parse_args()

# Biopython's trees don't store links to node parents, so we need to build
# a map of each node to its parent.
# Code from the Bio.Phylo cookbook: http://biopython.org/wiki/Phylo_cookbook
def all_parents(tree):
    parents = {}
    for clade in tree.find_clades(order='level'):
        for child in clade:
            parents[child] = clade
    return parents

def annotate_phylotree_parents(tree):
    # Get all parent nodes by node.
    parents_by_node = all_parents(tree)

    # Next, annotate each node with its parent.
    for node in tree.find_clades():
        if node == tree.root:
            node.parent = None
        else:
            node.parent = parents_by_node[node]

    # Return the tree.
    return tree


def modified_tree_to_json(node, extra_attr = []):
    tree_json = {}
    str_attr = ['strain','attr']
    num_attr = ['yvalue', 'tvalue', 'num_date', 'clade']
    if hasattr(node, 'name'):
        tree_json['strain'] = node.name

    for prop in str_attr:
        if hasattr(node, prop):
            tree_json[prop] = node.__getattribute__(prop)
    for prop in num_attr:
        if hasattr(node, prop):
            try:
                tree_json[prop] = round(node.__getattribute__(prop),5)
            except:
                print("cannot round:", node.__getattribute__(prop), "assigned as is")
                tree_json[prop] = node.__getattribute__(prop)

    for prop in extra_attr:
        if len(prop)==2 and callable(prop[1]):
            if hasattr(node, prop[0]):
                tree_json[prop] = prop[1](node.__getattribute__(prop[0]))
        else:
            if hasattr(node, prop):
                tree_json[prop] = node.__getattribute__(prop)

    if node.clades:
        tree_json["children"] = []
        for ch in node.clades:
            tree_json["children"].append(modified_tree_to_json(ch, extra_attr))
    return tree_json

def mock_meta_json(tree, args):
    meta = {}
    meta["updated"] = "today"
    meta["virus_count"] = tree.virus_count
    meta["maintainer"] = args.maintainer

    if not args.title:
        meta["title"] = args.newick
    else:
        meta["title"] = args.title

    meta["color_options"] = {
        # "country": {
        #     "menuItem": "country",
        #     "type": "discrete",
        #     "legendTitle": "country",
        #     "key": "country"
        #   },
        "num_date": {
            "menuItem": "date",
            "type": "continuous",
            "legendTitle": "Sampling date",
            "key": "num_date"
        },
    }
    meta["filters"] = []
    for filter in args.filters:
        meta["filters"].append(filter)
    meta["commit"] = "unknown"
    meta["panels"] = ["tree"]
    meta["geo"] = {"country": {}}
    meta["annotations"] = {}
    meta["author_info"] = {}
    if args.defaults:
        meta["defaults"] = {}
        for default in args.defaults:
            kv = default.split(":")
            meta["defaults"][kv[0]] = kv[1]


    return meta;

def set_basic_information_on_nodes(tree):
    count = 0;
    for node in tree.find_clades():
        count += 1
        if not node.name:
            node.name = "CLADE_{}".format(count)
        setattr(node,'attr', {})
        setattr(node,'clade', count)
        setattr(node,'strain', node.name)

def set_y_values(tree):
    count = 0;
    for node in tree.find_clades():
        if node.is_terminal():
            setattr(node, 'yvalue', count)
            count += 1
    # set internal y-values
    for node in tree.get_nonterminals(order="postorder"):
        setattr(node, 'yvalue', np.mean([x.yvalue for x in node.clades]))
    tree.virus_count = count

def set_divergence_on_tree(tree):
    for node in tree.find_clades():
        if node.name == "root":
            node.branch_length = 0.0 # root node
            node.cumulative_length = 0.0
        else:
            node.cumulative_length = node.parent.cumulative_length + node.branch_length
        node.attr["div"] = node.cumulative_length

def make_up_temporal_data(tree):

    fake_date_range = [2000, 2018]
    max_divergence = 0;
    for node in tree.find_clades():
        if node.attr["div"] > max_divergence:
            max_divergence = node.attr["div"]

    for node in tree.find_clades():
        node.attr["num_date"] = fake_date_range[0] + (node.attr["div"] / max_divergence) * (fake_date_range[1] - fake_date_range[0])


def make_up_country(tree):
    countries = ["USA", "New Zealand", "France", "Kenya"]
    for node in tree.find_clades():
        if node.is_terminal():
            node.attr["country"] = np.random.choice(countries)

def calc_entropy(vals):
    if len(vals) == 1:
        return 0.0
    return sum([v * math.log(v+1E-20) for v in vals]) * -1 / math.log(len(vals))


def removeBOM(arr):
    return [field.replace('\xef\xbb\xbf', '') for field in arr] # excel BOM \U+FEFF

if __name__=="__main__":
    args = get_command_line_args()
    root_logger = logging.getLogger('')
    root_logger.setLevel(args.loglevel if args.loglevel else logging.INFO)
    root_logger.addHandler(ColorizingStreamHandler())
    logger = logging.getLogger(__name__)


    ## PARSE THE TREE ##
    if args.newick:
        logger.info("Loading newick tree {}".format(args.newick))
        tree = Phylo.read(args.newick, "newick");
        tree.ladderize()
        tree = annotate_phylotree_parents(tree)
        set_basic_information_on_nodes(tree)
        set_y_values(tree)
        set_divergence_on_tree(tree)
        make_up_temporal_data(tree)
        make_up_country(tree)
    elif args.nexus:
        try:
            bt = imp.load_source('baltic', '/Users/james/blab/baltic/baltic.py')
        except IOError:
            bt = imp.load_source('baltic', '/Users/evogytis/Documents/BLAB_baltic/baltic.py')
        logger.info("Loading nexus tree {}".format(args.nexus))
        bt_tree = bt.loadNexus(args.nexus, absoluteTime=False) ## loads a BEAST nexus file
        bt_tree.setAbsoluteTime(args.most_recent_tip)

        tree = Phylo.read(StringIO(bt_tree.toString()), 'newick')
        tree = annotate_phylotree_parents(tree)
        nodes=tree.get_nonterminals(order='preorder') ## fetch nodes from newick tree
        bt_nodes=[k for k in bt_tree.traverse_tree(include_all=True) if k.branchType=='node'] ## fetch nodes from baltic tree
        leaves=tree.get_terminals(order='preorder') ## fetch leaves from newick tree
        bt_leaves=[k for k in bt_tree.traverse_tree()] ## fetch leaves from baltic tree
        tree.virus_count = len(leaves)

        # iterate over BT (nexus/BEAST) nodes & leaves as we iterate over the biophylo tree
        # and transfer relevent data
        for pairs in [[nodes,bt_nodes],[leaves,bt_leaves]]:
            for node, b in zip(*pairs):
                setattr(node,'tvalue',b.height) ## height in years
                setattr(node,'xvalue',b.height) ## 0.0 in the future?
                setattr(node,'yvalue',b.y) ## y value after drawing
                setattr(node,'clade',b.index) ## clade

                attrs={}

                if b.absoluteTime!=None:
                    attrs['num_date']=b.absoluteTime
                else: # root
                    attrs['num_date']=bt_tree.Objects[0].absoluteTime

                if b.traits.has_key('height_95%_HPD'):
                    attrs['num_date_confidence']=[args.most_recent_tip-t for t in b.traits['height_95%_HPD']]


                if args.fake_divergence:
                    attrs['div'] = attrs['num_date']


                setattr(node,'attr',attrs)

                for trait in args.discrete_traits:
                    if b.traits.has_key(trait):
                        attrs[trait] = b.traits[trait]
                        tset = trait+".set"
                        tprob = trait+".set.prob"
                        if b.traits.has_key(tset) and b.traits.has_key(tprob):
                            attrs[trait+'_confidence']={t:p for t,p in zip(b.traits[tset], b.traits[tprob])}
                            attrs[trait+'_entropy'] = calc_entropy(b.traits[tprob])

                for trait in args.continuous_traits:
                    make_log = trait in args.make_traits_log
                    rate = trait+".rate"
                    hpd = trait+".rate_95%_HPD"
                    if b.traits.has_key(rate):
                        attrs[trait] = np.log(b.traits[rate]) if make_log else b.traits[rate]
                        if b.traits.has_key(hpd):
                            attrs[trait+'_confidence'] = [np.log(x) for x in b.traits[hpd]] if make_log else b.traits[hpd]
                            attrs[trait+'_entropy'] = calc_entropy(b.traits[hpd])


        # import pdb; pdb.set_trace()
    meta_json = mock_meta_json(tree, args)


    if args.strain_csv:
        meta = {}
        with open(args.strain_csv, 'rb') as fh:
            csvdata = csv.reader(fh, delimiter=',', quotechar='"')
            header = removeBOM(csvdata.next())
            assert(header[0] == "strain")
            for line in csvdata:
                meta[line[0]] = {k:v for k,v in zip(header, line) if k != "strain"}

        for node in tree.find_clades():
            if node.name in meta:
                if not hasattr(node, "attr"): # confusing!
                    setattr(node, "attr", {})
                for key, value in meta[node.name].items():
                    if value == "":
                        continue
                    node.attr[key] = value

        for trait in header[1:]:
            meta_json["color_options"][trait] = {"menuItem": trait, "type": "discrete", "legendTitle": trait, "key": trait}

    if args.geo:
        geo = defaultdict(lambda: defaultdict(dict))
        with open(args.geo, 'rb') as fh:
            csvdata = csv.reader(fh, delimiter=',', quotechar='"')
            header = removeBOM(csvdata.next())
            assert(len(header) == 4)
            for line in csvdata:
                # trait,value,latitude,longitude
                geo[line[0]][line[1]] = {"latitude": line[2], "longitude": line[3]}
        meta_json["geo"] = geo
        if "map" not in meta_json["panels"]:
            meta_json["panels"].append("map")

        # some basic checks
        geo_traits = geo.keys()
        for node in tree.find_clades():
            if hasattr(node, "attr"):
                for trait in geo_traits:
                    if trait in node.attr:
                        if node.attr[trait] not in geo[trait]:
                            logger.warn("{} has no GPS co-ords".format(node.attr[trait]))

    tree_json = modified_tree_to_json(tree.root)

    if args.nexus:
        tree_json = tree_json["children"][0] # hacky

    ## add the discrete & continuous traits as colorBys
    for trait in args.discrete_traits:
        meta_json["color_options"][trait] = {"menuItem": trait, "type": "discrete", "legendTitle": trait, "key": trait}
    for trait in args.continuous_traits:
        meta_json["color_options"][trait] = {"menuItem": trait, "type": "continuous", "legendTitle": trait, "key": trait}

    json.dump(tree_json, open("{}_tree.json".format(args.output_prefix), 'w'), indent=2)
    json.dump(meta_json, open("{}_meta.json".format(args.output_prefix), 'w'), indent=2)

    logger.info("DONE")