#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Given an outputfile from one of the PHYLIP tools - `dnaml` or `dnapars` -
produce a CollapsedForest containing the trees in that outputfile."""

import gctree.branching_processes as bp
import gctree.utils

from ete3 import Tree
import re
import random
from collections import defaultdict

import pickle
import argparse
from warnings import warn
from typing import Dict


class TreeParseError(RuntimeError):
    pass


code_vectors = {
    code: [
        0 if base in gctree.utils.ambiguous_dna_values[code] else float("inf")
        for base in gctree.utils.bases
    ]
    for code in gctree.utils.ambiguous_dna_values
}
cost_adjust = {
    base: [int(not i == j) for j in range(5)]
    for i, base in enumerate(gctree.utils.bases)
}


# iterate over recognized sections in the phylip output file.
def sections(fh):
    patterns = {
        "parents": r"\s+between\s+and\s+length",
        ("sequences", "dnaml"): r"\s*node\s+reconstructed\s+sequence",
        ("sequences", "dnapars"): r"from\s+to\s+any steps",
        "seqboot_dataset": r"Data\s*set",
    }
    patterns = {k: re.compile(v, re.IGNORECASE) for (k, v) in patterns.items()}
    for line in fh:
        for k, pat in patterns.items():
            if pat.match(line):
                yield k
                break


# iterate over entries in the distance section
def iter_edges(fh):
    #  152          >root2           0.01208     (     zero,     0.02525) **
    pat = re.compile(
        r"\s*(?P<parent>\w+)\s+(?P<child>[\w>_.-]+)\s+(?P<distance>\d+\.\d+)"
    )
    # drop the header underline
    fh.readline()
    matches = 0
    for line in fh:
        m = pat.match(line)
        if m:
            matches += 1
            yield (m.group("child"), m.group("parent"))
        # We only want to break at the very end of the block of matches; dnaml has an extra blank line between
        # header and rows that dnapars doesn't
        elif not matches:
            continue
        else:
            break


# iterate over entries in the sequences section
def parse_seqdict(fh, mode="dnaml"):
    #  152        sssssssssG AGGTGCAGCT GTTGGAGTCT GGGGGAGGCT TGGTACAGCC TGGGGGGTCC
    seqs = defaultdict(str)
    if mode == "dnaml":
        patterns = re.compile(r"^\s*(?P<id>[a-zA-Z0-9>_.-]*)\s+(?P<seq>[a-zA-Z \-]+)")
    elif mode == "dnapars":
        patterns = re.compile(
            r"^\s*\S+\s+(?P<id>[a-zA-Z0-9>_.-]*)\s+(yes\s+|no\s+|maybe\s+)?(?P<seq>[a-zA-Z \-\?]+)"
        )
    else:
        raise ValueError("invalid mode " + mode)
    fh.readline()
    for line in fh:
        m = patterns.match(line)
        if m and m.group("id") != "":
            last_blank = False
            seqs[m.group("id")] += m.group("seq").replace(" ", "").upper()
        elif line.rstrip() == "":
            if last_blank:
                break
            else:
                last_blank = True
                continue
        else:
            break
    return seqs


def get_expected_trees(outfile):
    with open(outfile) as fh:
        for line, _ in zip(fh, range(7)):
            if "trees in all found" in line:
                return int(line.strip().split()[0])
                break
        else:
            return None


# parse the dnaml output file and return data structures containing a
# list biopython.SeqRecords and a dict containing adjacency
# relationships and distances between nodes.
def parse_outfile(outfile, abundance_file=None, root="root", disambiguate=False):
    """Parse phylip outfile, and return dnapars trees."""
    _expected_seq_len = None
    _expected_n_trees = get_expected_trees(outfile)
    _n_trees_warned = False
    if abundance_file is not None:
        counts = {
            line.split(",")[0]: int(line.split(",")[1]) for line in open(abundance_file)
        }
    # No count, just make an empty count dictionary:
    else:
        counts = None
    trees = []
    bootstrap = False
    # Ugg... for compilation need to let python know that these will definely both be defined :-/
    sequences, parents = {}, {}
    with open(outfile, "r") as fh:
        for sect in sections(fh):
            try:
                if sect == "parents":
                    parents = {child: parent for child, parent in iter_edges(fh)}
                elif sect[0] == "sequences":
                    sequences = parse_seqdict(fh, sect[1])
                    # sanity check;  a valid tree should have exactly one node that is parentless
                    if not len(parents) == len(sequences) - 1:
                        raise TreeParseError(
                            "invalid results attempting to parse {}: there are {} parentless sequences".format(
                                outfile, len(sequences) - len(parents)
                            )
                        )
                    # Also, all sequences should have the same length!
                    _seq_lengths = list(set(len(seq) for _, seq in sequences.items()))
                    if not len(_seq_lengths) == 1:
                        raise TreeParseError(
                            "Parsed node sequences do not all have the same length!"
                        )
                    # Also, all sequences' lengths should match first tree's
                    # seq lengths
                    if _expected_seq_len is not None:
                        if not _seq_lengths[0] == _expected_seq_len:
                            raise TreeParseError(
                                "Parsed tree node sequences don't match first parsed tree's sequences!"
                            )
                    else:
                        _expected_seq_len = _seq_lengths[0]

                    if bootstrap:
                        trees[-1].append(build_tree(sequences, parents, counts, root))
                    else:
                        trees.append(build_tree(sequences, parents, counts, root))
                elif sect == "seqboot_dataset":
                    bootstrap = True
                    trees.append([])
                else:
                    raise TreeParseError(
                        "unrecognized phylip section = {}".format(sect)
                    )
            except TreeParseError:
                if _expected_n_trees is not None:
                    expected_message = f" but {_expected_n_trees} were expected"
                else:
                    expected_message = ""
                warn(
                    f"Error parsing {outfile}! {len(trees)} successfully read{expected_message}."
                )
                _n_trees_warned = True
                break
        if (
            not _n_trees_warned
            and _expected_n_trees is not None
            and _expected_n_trees != len(trees)
        ):
            warn(
                f"Parsed {len(trees)} from {outfile} but {_expected_n_trees} expected!"
            )
    if len(trees) == 0:
        raise RuntimeError(f"No trees found in '{outfile}'")
    if disambiguate:
        # Disambiguate sets node.dist for all nodes in disambiguated trees
        trees = [disambiguate(tree) for tree in trees]
    return trees


def disambiguate(tree: Tree, random_state=None) -> Tree:
    """Randomly resolve ambiguous bases using a two-pass Sankoff Algorithm on
    subtrees of consecutive ambiguity codes."""
    if random_state is None:
        random.seed(tree.write(format=1))
    else:
        random.setstate(random_state)
    for node in tree.traverse():
        for site, base in enumerate(node.sequence):
            if base not in gctree.utils.bases:

                def is_leaf(node):
                    return (node.is_leaf()) or (
                        node.sequence[site] in gctree.utils.bases
                    )

                # First pass of Sankoff: compute cost vectors
                for node2 in node.traverse(strategy="postorder", is_leaf_fn=is_leaf):
                    base2 = node2.sequence[site]
                    node2.add_feature("cv", code_vectors[base2].copy())
                    if not is_leaf(node2):
                        for i in range(5):
                            for child in node2.children:
                                node2.cv[i] += min(
                                    [
                                        sum(v)
                                        for v in zip(
                                            child.cv, cost_adjust[gctree.utils.bases[i]]
                                        )
                                    ]
                                )
                # Second pass: Choose base and adjust children's cost vectors
                if not node.is_root():
                    node.cv = [
                        sum(v)
                        for v in zip(node.cv, cost_adjust[node.up.sequence[site]])
                    ]
                # traverse evaluates is_leaf(node) after yielding node.
                # Resolving base makes is_leaf true; must get order before
                # making changes.
                preorder = list(node.traverse(strategy="preorder", is_leaf_fn=is_leaf))
                for node2 in preorder:
                    if node2.sequence[site] in gctree.utils.bases:
                        continue
                    min_cost = min(node2.cv)
                    base_index = random.choice(
                        [i for i, val in enumerate(node2.cv) if val == min_cost]
                    )
                    new_base = gctree.utils.bases[base_index]
                    # Adjust child cost vectors
                    if not is_leaf(node2):
                        for child in node2.children:
                            child.cv = [
                                sum(v) for v in zip(child.cv, cost_adjust[new_base])
                            ]
                    node2.sequence = (
                        node2.sequence[:site] + new_base + node2.sequence[(site + 1) :]
                    )
    for node in tree.traverse():
        try:
            node.del_feature("cv")
        except (AttributeError, KeyError):
            pass
    tree.dist = 0
    for node in tree.iter_descendants():
        node.dist = gctree.utils.hamming_distance(node.up.sequence, node.sequence)
    return tree


# build a tree from a set of sequences and an adjacency dict.
def build_tree(
    sequences: Dict[str, str], parents: Dict[str, str], counts=None, root="root"
):
    """Build an ete tree from sequences and parents dictionaries.

    Args:
        sequences: a dictionary mapping names to sequences
        parents: a dictionary mapping parent node names to child node names
        counts: a dictionary mapping node names to observed abundances. This argument
            is no longer used in the main gctree inference pipeline (counts are assigned in DAG)
            but remains for compatibility.
        root: the name of the root node
    """
    # build an ete tree
    # first a dictionary of disconnected nodes
    nodes = {}
    for name in sequences:
        node = Tree()
        node.name = name
        node.add_feature("sequence", sequences[node.name])
        if counts is not None:
            if node.name in counts:
                node.add_feature("abundance", counts[node.name])
            else:
                node.add_feature("abundance", 0)
        nodes[name] = node
    for name in sequences:
        if name in parents:
            nodes[parents[name]].add_child(nodes[name])
        else:
            tree = nodes[name]
    # reroot on root
    if root is not None:
        try:
            root_id = [node for node in nodes if root in node][0]
        except IndexError:
            raise RuntimeError(f"Provided root id '{root}' not found in dnapars tree.")
        assert len(nodes[root_id].children) == 0
        root_parent = nodes[root_id].up
        root_parent.remove_child(nodes[root_id])
        nodes[root_id].add_child(root_parent)
        # remove possible unecessary unifurcation after rerooting
        if len(root_parent.children) == 1:
            root_parent.delete(prevent_nondicotomic=False)
            root_parent.children[0].dist = gctree.utils.hamming_distance(
                root_parent.children[0].sequence, nodes[root_id].sequence
            )
        tree = nodes[root_id]

    return tree


def get_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "phylip_outfile",
        type=str,
        help="dnaml outfile (verbose output with inferred ancestral sequences, option 5).",
    )
    parser.add_argument("abundance_file", type=str, help="count file")
    parser.add_argument("--outputfile", default=None, help="output file.")
    parser.add_argument("--root", default="root", help="root sequence id")
    return parser


def main(arg_list=None):
    args = get_parser().parse_args(arg_list)

    if args.outputfile is None:
        args.outputfile = args.phylip_outfile + ".collapsed_forest.p"
    forest = bp.CollapsedForest(
        parse_outfile(args.phylip_outfile, args.abundance_file, args.root)
    )
    with open(args.outputfile, "wb") as fh:
        pickle.dump(forest, fh)


if __name__ == "__main__":
    main()
