#!/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 collections
import os
import random

from paleomix.common.fileutils import add_postfix, swap_ext
from paleomix.nodes.examl import ExaMLNode, ExaMLParserNode
from paleomix.nodes.formats import FastaToPartitionedInterleavedPhyNode as ToPhylipNode
from paleomix.nodes.newick import NewickRerootNode, NewickSupportNode
from paleomix.nodes.phylip import PHYLIPBootstrapNode
from paleomix.nodes.raxml import RAxMLParsimonyTreeNode


def _build_supermatrix(
    destination, input_files, exclude_samples, subset_files, dependencies
):
    matrixprefix = os.path.join(destination, "alignments")

    # This supermatrix (and partitions) is not used by the pipeline, but is
    # built for reference purposes; may be made optional in the future.
    supermatrix = ToPhylipNode(
        infiles=input_files,
        out_prefix=matrixprefix,
        exclude_groups=exclude_samples,
        dependencies=dependencies,
        file_dependencies=subset_files,
    )

    # Supermatrix with empty columns (all N, n, or -) stripped
    reduced = ToPhylipNode(
        reduce=True,
        infiles=input_files,
        out_prefix=matrixprefix + ".reduced",
        exclude_groups=exclude_samples,
        dependencies=dependencies,
        file_dependencies=subset_files,
    )

    return (supermatrix, reduced)


def _examl_nodes(
    options,
    settings,
    input_alignment,
    input_partitions,
    input_binary,
    output_template,
    dependencies,
):
    parsimony_tree = output_template % ("parsimony_tree",)
    tree = RAxMLParsimonyTreeNode(
        input_alignment=input_alignment,
        input_partitions=input_partitions,
        output_tree=parsimony_tree,
        dependencies=dependencies,
    )

    return ExaMLNode(
        input_binary=input_binary,
        initial_tree=parsimony_tree,
        output_template=output_template,
        model=settings["ExaML"]["Model"].upper(),
        threads=options.examl_max_threads,
        dependencies=tree,
    )


def _build_rerooted_trees(nodes, reroot_on):
    filenames = []
    for node in nodes:
        for filename in node.output_files:
            if filename.endswith(".result"):
                filenames.append(filename)

    output_file = os.path.dirname(filenames[0]) + ".newick"

    return NewickRerootNode(
        tree_files=filenames,
        output_file=output_file,
        taxa=reroot_on,
        dependencies=nodes,
    )


def _build_examl_replicates(
    options, phylo, destination, input_alignment, input_partition, dependencies
):
    input_binary = os.path.join(destination, "alignments.reduced.binary")
    binary = ExaMLParserNode(
        input_alignment=input_alignment,
        input_partition=input_partition,
        output_file=input_binary,
        dependencies=dependencies,
    )

    replicates = []
    for replicate_num in range(phylo["ExaML"]["Replicates"]):
        replicate_destination = os.path.join(destination, "replicates")
        replicate_template = os.path.join(
            replicate_destination, "replicate.%04i.%%s" % (replicate_num,)
        )
        replicates.append(
            _examl_nodes(
                options,
                phylo,
                input_alignment,
                input_partition,
                input_binary,
                replicate_template,
                binary,
            )
        )

    if replicates:
        return _build_rerooted_trees(replicates, phylo["RootTreesOn"])

    return None


def _build_examl_bootstraps(
    options, phylo, destination, input_alignment, input_partition, dependencies
):
    bootstraps = []
    num_bootstraps = phylo["ExaML"]["Bootstraps"]
    bootstrap_destination = os.path.join(destination, "bootstraps")
    bootstrap_template = os.path.join(bootstrap_destination, "bootstrap.%04i.phy")

    for bootstrap_num in range(num_bootstraps):
        bootstrap_alignment = bootstrap_template % (bootstrap_num,)
        bootstrap = PHYLIPBootstrapNode(
            input_alignment=input_alignment,
            input_partition=input_partition,
            output_alignment=bootstrap_alignment,
            seed=random.randint(1, 2**32 - 1),
            dependencies=dependencies,
        )

        bootstrap_binary = swap_ext(bootstrap_alignment, ".binary")
        bootstrap_final = swap_ext(bootstrap_alignment, ".%s")
        bs_binary = ExaMLParserNode(
            input_alignment=bootstrap_alignment,
            input_partition=input_partition,
            output_file=bootstrap_binary,
            dependencies=bootstrap,
        )

        bootstraps.append(
            _examl_nodes(
                options=options,
                settings=phylo,
                input_alignment=bootstrap_alignment,
                input_partitions=input_partition,
                input_binary=bootstrap_binary,
                output_template=bootstrap_final,
                dependencies=bs_binary,
            )
        )

    if bootstraps:
        return _build_rerooted_trees(bootstraps, phylo["RootTreesOn"])

    return None


def add_bootstrap_support(destination, replicate, bootstrap):
    if not (replicate and bootstrap):
        return [_f for _f in (replicate, bootstrap) if _f]

    replicate_file = os.path.join(destination, "replicates.newick")
    bootstrap_file = os.path.join(destination, "bootstraps.newick")
    output_file = add_postfix(replicate_file, ".support")

    return (
        NewickSupportNode(
            main_tree_files=replicate_file,
            support_tree_files=bootstrap_file,
            output_file=output_file,
            dependencies=(bootstrap, replicate),
        ),
    )


def _build_examl_nodes(
    options, settings, destination, input_files, subset_files, dependencies
):
    input_alignment = os.path.join(destination, "alignments.reduced.phy")
    input_partition = os.path.join(destination, "alignments.reduced.partitions")

    excluded = settings["ExcludeSamples"]
    supermatrix = _build_supermatrix(
        destination, input_files, excluded, subset_files, dependencies
    )
    examl_args = (
        options,
        settings,
        destination,
        input_alignment,
        input_partition,
        supermatrix,
    )

    examl_replicates = _build_examl_replicates(*examl_args)
    examl_bootstraps = _build_examl_bootstraps(*examl_args)
    examl_dependencies = add_bootstrap_support(
        destination, examl_replicates, examl_bootstraps
    )

    return examl_dependencies


def _build_examl_per_gene_nodes(
    options, settings, run_dd, roi, destination, filtering, dependencies
):
    regions = settings["Project"]["Regions"][roi["Name"]]
    sequences = regions["Sequences"][roi["SubsetRegions"]]
    subset_files = regions["SubsetFiles"][roi["SubsetRegions"]]
    filtering_postfix = ".filtered" if any(filtering.values()) else ""
    sequence_dir = os.path.join(
        options.destination, "alignments", roi["Name"] + filtering_postfix
    )
    msa_enabled = settings["MultipleSequenceAlignment"][regions["Name"]]["Enabled"]
    fasta_extension = ".afa" if msa_enabled else ".fasta"

    partitions = roi["Partitions"] or "111"

    nodes = []
    for sequence in sequences:
        seq_source = os.path.join(sequence_dir, sequence + fasta_extension)
        seq_destination = os.path.join(destination, sequence)
        input_files = {sequence: {"partitions": partitions, "filenames": [seq_source]}}
        nodes.extend(
            _build_examl_nodes(
                options,
                run_dd,
                seq_destination,
                input_files,
                subset_files,
                dependencies,
            )
        )

    return nodes


def _build_examl_regions_nodes(
    options, settings, run_dd, destination, filtering, dependencies
):
    input_files = collections.defaultdict(dict)
    subset_files = []
    for roi_name, roi_dd in run_dd["RegionsOfInterest"].items():
        regions = settings["Project"]["Regions"][roi_name]
        subset_key = roi_dd.get("SubsetRegions")
        sequences = regions["Sequences"][subset_key]
        subset_files.extend(regions["SubsetFiles"][subset_key])

        partitions = roi_dd["Partitions"]
        filtering_postfix = ".filtered" if any(filtering.values()) else ""
        sequence_dir = os.path.join(
            options.destination, "alignments", roi_name + filtering_postfix
        )
        msa_enabled = settings["MultipleSequenceAlignment"][regions["Name"]]["Enabled"]
        fasta_extension = ".afa" if msa_enabled else ".fasta"

        if partitions:
            for sequence in sequences:
                seq_source = os.path.join(sequence_dir, sequence + fasta_extension)

                input_files[sequence] = {
                    "partitions": partitions,
                    "filenames": [seq_source],
                }
        else:
            filenames = []
            for sequence in sequences:
                filenames.append(os.path.join(sequence_dir, sequence + fasta_extension))

            input_files[roi_name] = {"partitions": "1", "filenames": filenames}

    return _build_examl_nodes(
        options, run_dd, destination, dict(input_files), subset_files, dependencies
    )


def build_phylogeny_nodes(options, settings, filtering, dependencies):
    nodes = []
    for run_name, run_dd in settings["PhylogeneticInference"].items():
        destination = os.path.join(options.destination, "phylogenies", run_name)

        if run_dd["PerGeneTrees"]:
            run_nodes = []
            for roi in run_dd["RegionsOfInterest"].values():
                roi_destination = os.path.join(destination, roi["Name"])
                run_nodes.extend(
                    _build_examl_per_gene_nodes(
                        options,
                        settings,
                        run_dd,
                        roi,
                        roi_destination,
                        filtering,
                        dependencies,
                    )
                )
            nodes.extend(run_nodes)
        else:
            nodes.extend(
                _build_examl_regions_nodes(
                    options, settings, run_dd, destination, filtering, dependencies
                )
            )

    return nodes


def chain_examl(_pipeline, options, makefiles):
    destination = options.destination  # Move to makefile
    for makefile in makefiles:
        filtering = makefile["Project"]["FilterSingletons"]
        options.destination = os.path.join(destination, makefile["Project"]["Title"])

        makefile["Nodes"] = build_phylogeny_nodes(
            options, makefile, filtering, makefile["Nodes"]
        )
    options.destination = destination
