#!/usr/bin/python3
#
# Copyright (c) 2012 Mikkel Schubert <MikkelSch@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import os

from paleomix.common.formats.newick import Newick
from paleomix.common.utilities import safe_coerce_to_tuple
from paleomix.common.fileutils import describe_files, move_file
from paleomix.node import Node


class NewickRerootNode(Node):
    def __init__(self, tree_files, output_file, taxa=(), dependencies=()):
        self._output_file = output_file
        self._tree_files = safe_coerce_to_tuple(tree_files)
        self._reroot_on_taxa = safe_coerce_to_tuple(taxa)

        reroot_on = "midpoint"
        if self._reroot_on_taxa:
            reroot_on = repr("', '".join(sorted(self._reroot_on_taxa)))

        Node.__init__(
            self,
            description="rerooting newick tree %s on %s"
            % (describe_files(tree_files), reroot_on),
            input_files=self._tree_files,
            output_files=self._output_file,
            dependencies=dependencies,
        )

    def _run(self, _config, temp):
        lines = []
        for tree in _read_tree_files(self._tree_files):
            if self._reroot_on_taxa:
                rooted_tree = tree.reroot_on_taxa(self._reroot_on_taxa)
            else:
                rooted_tree = tree.reroot_on_midpoint()
            lines.append(str(rooted_tree))
        lines = "\n".join(lines) + "\n"

        temp_output_file = os.path.join(temp, os.path.basename(self._output_file))
        with open(temp_output_file, "w") as handle:
            handle.write(lines)

        move_file(temp_output_file, self._output_file)


class NewickSupportNode(Node):
    def __init__(
        self, main_tree_files, support_tree_files, output_file, dependencies=()
    ):
        self._output_file = output_file
        self._main_tree_files = safe_coerce_to_tuple(main_tree_files)
        self._support_tree_files = safe_coerce_to_tuple(support_tree_files)
        input_files = self._main_tree_files + self._support_tree_files

        Node.__init__(
            self,
            description="calculating newick support for %s"
            % (describe_files(main_tree_files),),
            input_files=input_files,
            output_files=output_file,
            dependencies=dependencies,
        )

    def _run(self, _config, temp):
        main_trees = _read_tree_files(self._main_tree_files)
        support_trees = _read_tree_files(self._support_tree_files)

        lines = []
        for main_tree in main_trees:
            supported_tree = main_tree.add_support(support_trees)
            lines.append(str(supported_tree))
        lines = "\n".join(lines) + "\n"

        temp_output_file = os.path.join(temp, os.path.basename(self._output_file))
        with open(temp_output_file, "w") as handle:
            handle.write(lines)

        move_file(temp_output_file, self._output_file)


def _read_tree_files(filenames):
    trees = []
    for filename in filenames:
        with open(filename) as handle:
            for line in handle:
                trees.append(Newick.from_string(line))
    return trees
