#!/usr/bin/env python
# ----------------------------------------------------------------------------
# Copyright (c) 2015--, ghost-tree development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the LICENSE file, distributed with this software.
# ----------------------------------------------------------------------------

# Avoid importing QtCore via skbio importing matplotlib by
# specifying AGG as matplotlib backend:
import matplotlib
matplotlib.use("AGG")

import click
import skbio
import os

# library code imports
from ghosttree.util import compare_tip_to_tip_distances
from ghosttree.filter import filter_positions
from ghosttree.extensions.cluster import preprocess_extension_tree_sequences
from ghosttree.silva.filter import fungi_from_fasta
from ghosttree.scaffold.hybridtree import extensions_onto_foundation


@click.group()
def cli():
    """\n\n
    ___________________________________________________________________

    ghost-tree: a tool for creating hybrid-gene phylogenetic trees for
    diversity analyses. (https://github.com/JTFouquier/ghost-tree)
    ___________________________________________________________________

    For dependent software installation instructions please see:
    https://github.com/JTFouquier/ghost-tree/blob/master/README.rst

    -> For help with each command, enter "--help"
    (e.g. 'ghost-tree scaffold hybrid-tree --help').

    -> For a more detailed tutorial, please see:
    https://github.com/JTFouquier/ghost-tree/tree/master/workflow

    -> There are three 'command groups' (silva, extensions, and scaffold),
    which include additional 'subcommands.' (e.g. enter 'ghost-tree silva' and
    you will see the command 'ghost-tree silva extract-fungi' listed)

    -> Arguments *must* be entered in order, without any "-"/dashes.
    ___________________________________________________________________
    """
    pass

# click command and args here are for the user
# cli group is the "invisible" top level group


@cli.command("compare-trees")
@click.argument("tree-file1", type=click.File("U"))
@click.argument("tree-file2", type=click.File("U"))
@click.option("-m", "--method", default="pearson",
              type=click.Choice(["pearson", "spearman"]),
              help="correlation method to use in Mantel test")
def compare_trees(tree_file1, tree_file2, method):
    """Creates distance matrices for two trees (tree-file1 and tree-file2) and
    finds the correlation between the trees using a Mantel test."""
    coeff, p_value, n = compare_tip_to_tip_distances(tree_file1,
                                                     tree_file2,
                                                     method)
    click.echo("Correlation coefficient: %f" % coeff)
    click.echo("p-value: %f" % p_value)
    click.echo("Number of overlapping tips: %d" % n)


@cli.command("filter-alignment-positions")
@click.argument("foundation-alignment-file", type=click.File("U"))
@click.argument("maximum-gap-frequency", type=float)
@click.argument("maximum-position-entropy", type=float)
@click.argument("output-alignment-file", type=click.STRING)
def filter_alignment_positions(foundation_alignment_file,
                               maximum_gap_frequency,
                               maximum_position_entropy,
                               output_alignment_file):
    """Filters an alignment file based
    on positional properties. This command takes a foundation-alignment-file
    and uses the maximum-gap-frequency (e.g. 0.50)
    and maximum-position-entropy (e.g. 0.80) arguments to filter positions
    and return a new alignment output-file. Typically users would want to
    remove high gap positions and high entropy positions."""
    skbio.write(filter_positions(foundation_alignment_file,
                                 maximum_gap_frequency,
                                 maximum_position_entropy),
                into=output_alignment_file,
                format="fasta")


# silva is nested under cli group (silva shows up on command line
# interface but cli does not
@cli.group()
def silva():
    """Command group for SILVA rRNA database alignment files"""
    pass


@silva.command("extract-fungi")
@click.argument("aligned-fasta-file", type=click.File("U"))
@click.argument("accession-file", type=click.File("U"))
@click.argument("taxonomy-file", type=click.File("U"))
@click.argument("output-file", type=click.STRING)
def extract_fungi(aligned_fasta_file, accession_file, taxonomy_file,
                  output_file):
    """'extract-fungi' uses files that are unique to the
    SILVA rRNA database project. http://www.arb-silva.de/.
    Because SILVA database provides information for many organisms,
    this script is designed to extract the fungi from a large database.

    Arguments:

    1) aligned-fasta-file:
    SILVA sequences. Each sequence identifier must be an accession number.

    2) accession-file:
    A tab-separated file mapping accession numbers to a mapping number in
    `taxonomy_map`. This file should contain exactly two columns:
    accession number and mapping number.

    3) taxonomy_file:
    A tab-separated file that identifes the taxonomy and rank of a mapping
    number in `accession_fh`. This file should contain exactly five
    columns beginning with taxonomy, mapping number and rank. The last two
    columns are ignored.

    4) output-file:
    The resulting SILVA alignment file to use for the foundation of the
    ghost-tree, if SILVA is your selected foundation database.

    """
    skbio.write(fungi_from_fasta(aligned_fasta_file, accession_file,
                taxonomy_file), into=output_file, format="fasta")


@cli.group()
def extensions():
    """Command group for preprocessing 'extension trees'"""
    pass


@extensions.command("group-extensions")
@click.argument("species-level-sequence-file", type=str)
@click.argument("similarity-threshold", type=float)
@click.argument("output-file", type=str)
def group_extensions(species_level_sequence_file, similarity_threshold,
                     output_file):
    """Creates an OTU map from a sequence file in .fasta format using SUMACLUST
    (http://metabarcoding.org/sumatra). Sequence file can be from a database or
    from representative sequences specific to one's own dataset. OTU clusters
    will be created according to the threshold. The lower the threshold, the
    more "unidentified" sequences will be accounted for because they will be
    grouped with other identified fungi, but this will create
    larger "extension tree" groups that form less accurate multiple sequence
    alignments. Changing the threshold value to a number drastically different
    from the standard 97 or 99 OTU similarity will alter the resulting tree
    improving or reducing its utility in diversity analyses and should be
    performed with caution.' """
    preprocess_extension_tree_sequences(species_level_sequence_file,
                                        similarity_threshold,
                                        output_file)


@cli.group()
def scaffold():
    """Command group for creating hybrid-gene trees."""
    pass


@scaffold.command("hybrid-tree-foundation-alignment")
@click.argument("extension-trees-otu-map", type=click.File("U"))
@click.argument("extension-trees-taxonomy-file", type=click.File("U"))
@click.argument("extension-trees-sequence-file", type=click.File("U"))
@click.argument("foundation-alignment-file", type=click.File("U"))
@click.argument("ghost-tree-output-folder", type=click.STRING)
@click.option('--graft-level', type=click.Choice(['p', 'c', 'o', 'f', 'g']),
              default='g',
              help="Taxonomic level to graft extensions to foundation. "
                   "Default is g (genus). p=phylum, c=class, o=order, "
                   "f=family.")
@click.option("--stderr", is_flag=True, help="Saves error log file "
                                             "(optional)")
@click.option("--save-foundation-alignment", is_flag=True,
              help="Saves non-redundant foundation alignment file (optional)")
@click.option("--save-foundation-tree", is_flag=True,
              help="Saves foundation phylogenetic tree file (optional)")
@click.option("--exclude-id-list", is_flag=True,
              help="Will not save accession id file (optional)")
def hybrid_tree(graft_level, stderr,
                save_foundation_alignment, save_foundation_tree,
                exclude_id_list, extension_trees_otu_map,
                extension_trees_taxonomy_file, extension_trees_sequence_file,
                foundation_alignment_file, ghost_tree_output_folder):
    """Make hybrid-tree using foundation alignment.

    Combines two genetic databases into one phylogenetic tree and is the
    final command in ghost-tree.

    Some genetic databases provide finer taxonomic resolution, but high
    sequence variability causes poor multiple sequence alignments
    (we consider these "extension trees").

    Other databases provide high quality phylogenetic information (hence they
    can be used as the "foundation" tree), but they provide poor quality
    taxonomic resolution.

    The output file is one phylogenetic tree in .nwk format, which has the
    benefits of both databases. This allows sequencing to be performed using
    only one primer set, the "extension trees" primer set.

    Arguments:

    1) extension-trees-otu-map:
    Tab-delimited text file containing OTU clusters in rows containing
    accession numbers only. Format can be 1) where the accession number
    is in the first column with only one column or 2) it can contain
    accession numbers clustered in tab-delimited rows containing more
    accession numbers, which are part of that OTU cluster (as in output of
    "ghost-tree group-extensions"). This file refers to the "extension
    trees". File references to sequence reads or sample numbers/names are
    not valid here. This is NOT the same as an OTU .biom table.

    2) extension-trees-taxonomy-file:
    Tab-delimited text file related to "extension trees" wih the 1st
    column being an accession number (same accession numbers in
    extension-trees-otu-map and extension-trees-taxonomy-file) and the 2nd
    column is the taxonomy ranking in the following format:
    k__Fungi;p__Basidiomycota;c__Agaricomycetes;o__Sebacinales;
    f__Sebacinaceae;g__unidentified;s__Sebacina

    3) extension-trees-sequence-file:
    The .fasta formated sequences for the "extension trees" genetic
    dataset. Sequence identifiers are the accession numbers. These
    accession numbers are the same as in the extension-trees-otu-map and
    extension-trees-taxonomy-file.

    4) foundation-alignment-file:
    File containing aligned sequences from a genetic marker database in .fasta
    format. This file refers to the "foundation" of the ghost-tree.

    .fasta contains accession numbers *and* taxonomy labels.

    Note: if you would like to use a newick tree foundation, use the
    hybrid-tree-foundation-tree command, also in the 'scaffold' command group.

    5) ghost-tree-output-folder:
    Output folder contains files including:

    a) The Newick formatted ghost-tree, which is the final output of the
    ghost-tree tool. This is a phylogenetic tree designed for
    downstream diversity analyses.

    b) Accession IDs from the ghost-tree.nwk file that you can use for
    downstream analyses tools

    c) log error file (this is an optional file that you can have if you
    type '--stderr')
    """
    _, std_error = extensions_onto_foundation(extension_trees_otu_map,
                                              extension_trees_taxonomy_file,
                                              extension_trees_sequence_file,
                                              foundation_alignment_file,
                                              ghost_tree_output_folder,
                                              graft_level, None)

    _script_helper(stderr, save_foundation_alignment, save_foundation_tree,
                   exclude_id_list, ghost_tree_output_folder)


@scaffold.command("hybrid-tree-foundation-tree")
@click.argument("extension-trees-otu-map", type=click.File("U"))
@click.argument("extension-trees-taxonomy-file", type=click.File("U"))
@click.argument("extension-trees-sequence-file", type=click.File("U"))
@click.argument("foundation-tree-file", type=click.File("U"))
@click.argument('foundation-taxonomy-file', type=click.File('U'))
@click.argument("ghost-tree-output-folder", type=click.STRING)
@click.option('--graft-level', type=click.Choice(['p', 'c', 'o', 'f', 'g']),
              default='g',
              help="Taxonomic level to graft extensions to foundation. "
                   "Default is g (genus). p=phylum, c=class, o=order, "
                   "f=family.")
@click.option("--stderr", is_flag=True, help="Saves error log file "
                                             "(optional)")
@click.option("--exclude-id-list", is_flag=True,
              help="Will not save accession id file (optional)")
def hybrid_tree(graft_level, foundation_taxonomy_file, stderr,
                exclude_id_list, extension_trees_otu_map,
                extension_trees_taxonomy_file, extension_trees_sequence_file,
                foundation_tree_file, ghost_tree_output_folder):
    """Make hybrid-tree using foundation tree. Combines two genetic
    databases into one phylogenetic tree using a phylogenetic tree
    (newick format) as the foundation and is the final command in ghost-tree.

    Some genetic databases provide finer taxonomic resolution, but high
    sequence variability causes poor multiple sequence alignments
    (we consider these "extension trees").

    Other databases provide high quality phylogenetic information (hence they
    can be used as the "foundation" tree), but they provide poor quality
    taxonomic resolution.

    The output file is one phylogenetic tree in .nwk format, which has the
    benefits of both databases. This allows sequencing to be performed using
    only one primer set, the "extension trees" primer set.

    Arguments:

    1) extension-trees-otu-map:
    Tab-delimited text file containing OTU clusters in rows containing
    accession numbers only. Format can be 1) where the accession number
    is in the first column with only one column or 2) it can contain
    accession numbers clustered in tab-delimited rows containing more
    accession numbers, which are part of that OTU cluster (as in output of
    "ghost-tree group-extensions"). This file refers to the "extension
    trees". File references to sequence reads or sample numbers/names are
    not valid here. This is NOT the same as an OTU .biom table.

    2) extension-trees-taxonomy-file:
    Tab-delimited text file related to "extension trees" wih the 1st
    column being an accession number (same accession numbers in
    extension-trees-otu-map and extension-trees-taxonomy-file) and the 2nd
    column is the taxonomy ranking in the following format:
    k__Fungi;p__Basidiomycota;c__Agaricomycetes;o__Sebacinales;
    f__Sebacinaceae;g__unidentified;s__Sebacina

    3) extension-trees-sequence-file:
    The .fasta formated sequences for the "extension trees" genetic
    dataset. Sequence identifiers are the accession numbers. These
    accession numbers are the same as in the extension-trees-otu-map and
    extension-trees-taxonomy-file.

    4) foundation-tree-file:
    .nwk tree is a tree with accession numbers. This file refers to the
    "foundation" of the ghost-tree.

    5) foundation-taxonomy-file:
    Tab separated taxonomy file containing accession IDs and taxonomy for
    the foundation phylogenetic tree (newick format).

    6) ghost-tree-output-folder:
    Output folder contains files including:

    a) The Newick formatted ghost-tree, which is the final output of the
    ghost-tree tool. This is a phylogenetic tree designed for
    downstream diversity analyses.

    b) Accession IDs from the ghost-tree.nwk file that you can use for
    downstream analyses tools

    c) log error file (this is an optional file that you can have if you
    type '--stderr')
    """
    _, std_error = extensions_onto_foundation(extension_trees_otu_map,
                                              extension_trees_taxonomy_file,
                                              extension_trees_sequence_file,
                                              foundation_tree_file,
                                              ghost_tree_output_folder,
                                              graft_level,
                                              foundation_taxonomy_file)
    save_foundation_alignment = None
    save_foundation_tree = None
    _script_helper(stderr, save_foundation_alignment, save_foundation_tree,
                   exclude_id_list, ghost_tree_output_folder)


def _script_helper(stderr, save_foundation_alignment, save_foundation_tree,
                   exclude_id_list, ghost_tree_output_folder):

    if stderr:
        logfile = open(ghost_tree_output_folder + "/ghost_tree_log.txt", "w")
        logfile.write(std_error)
        logfile.close()

    if save_foundation_alignment:
        pass
    else:
        try:
            os.unlink(ghost_tree_output_folder +
                      "/nr_foundation_alignment_gt.fasta")
        except FileNotFoundError:
            pass

    if save_foundation_tree:
        pass
    else:
        try:
            os.unlink(ghost_tree_output_folder + "/nr_foundation_tree_gt.nwk")
        except FileNotFoundError:
            pass

    if exclude_id_list:
        os.unlink(ghost_tree_output_folder +
                  "/ghost_tree_extension_accession_ids.txt")


if __name__ == "__main__":
    cli()
