ct

#!/usr/bin/env python3

#                      _              _     
#                     | |            | |    
#   ___ __ _  __ _ ___| |_ ___   ___ | |___ 
#  / __/ _` |/ _` / __| __/ _ \ / _ \| / __|
# | (_| (_| | (_| \__ \ || (_) | (_) | \__ \
#  \___\__,_|\__,_|___/\__\___/ \___/|_|___/


'''

A Convergent Amino Acid Substitution identification 
and analysis toolbox

Author:         Fabio Barteri (fabio.barteri@upf.edu)

Contributors:   Alejandro Valenzuela (alejandro.valenzuela@upf.edu)
                Xavier Farré (xfarrer@igtp.cat),
                David de Juan (david.juan@upf.edu).

MODULE NAME:    ct
DESCRIPTION:    Launcher script. Calls the different tools.
DEPENDENCIES:   Directly or indirectly, all the modules are connected to this script

Version 0.9.1 - May 2023

'''

### Introduction

# Castools is a set of bioinformatics tools to detect and analyze
# convergent amino acid substitutions from multiple sequence alignments (MSA) of orthologous proteins.  

### Application info and general help

application_info = '''
CAASTOOLS version 0.9 (beta)
Convergent Amino Acid Substitution detection and analysis TOOLbox
'''

genhelp = '''
General usage:          > ct [tool] [options]
Help for single tool:   > ct [tool] --help

Tools           Description
--------        -------------------------------------------------

discovery       Detects Convergent Amino Acid Substitutions (CAAS) from
                a single Multiple Sequence Alignment (MSA).

resample        Resamples virtual phenotypes for CAAS bootstrap analysis.

bootstrap       Runs CAAS bootstrap analysis on a on a single MSA.

'''

### Imports

import sys                                          # for tool selection and more
from optparse import OptionParser                   # for single option parser
from modules.deps import check_dependencies         # 
from os.path import exists




#### PROGRAM CHECK ####################################################################################################
########################################################################################################################


if len(sys.argv) == 1:                                                  # The user didn't tape anything after "ct"
    print(application_info)
    print(genhelp)                                                      # Print toolbox-wide help
    print("\n\nPlease, select a tool\n\n")
    exit()

tool = sys.argv[1]

if "help" in tool.lower() or tool.lower == "-h":                                              # The user requested the toolbox-wide help prompt (genhelp)
    print(application_info)
    print(genhelp)                                                      # Print toolbox-wide help
    exit()

if tool.lower() not in ("discovery", "resample", "bootstrap"):          # Check: the user mistyped the name of a tool
    print(application_info)
    print(genhelp)                                                      # Print toolbox-wide help
    print("\n\n****ERROR: no tool named", tool + "\n\n")
    exit()




#### TOOL 1. DISCOVERY ######################################################################################################
########################################################################################################################

if tool.lower() == "discovery":

    ### 1.1 Check the dependencies
    check_dependencies("discovery", ["biopython", "scipy", "numpy"])

    ### 1.2 Init the input parser
    parser = OptionParser()


    ### 1.3 The mandatory inputs

    ###     1.3.1 Input alignment
    parser.add_option("-a", "--alignment", dest="single_alignment",
                    help="Multiple Sequence Alignment (MSA) file.", default = "none")
    ###     1.3.2 Alignment format
    parser.add_option("--fmt", dest="ali_format",
                    help="File format of the MSA file. Default: clustal. Accepted: clustal, emboss, fasta, \
                    fasta-m10, ig, maf, mauve, msf, nexus, phylip, phylip-sequential, phylip-relaxed, stockholm.", default = "clustal")
    ###     1.3.3 Config file
    parser.add_option("-t", "--traitfile", dest="config_file",
                    help="The trait config file (read documentation for file formatting)", default = "none")
    
    ###     1.3.4 Output file (the table)
    parser.add_option("-o", "--output", dest="output_file",
                    help="The output file, where the CAAS discovery table will be printed", default = "none")
    


    ### 1.4 Optional and filtering inputs

    ###     1.4.1 limit output to pattern (default: all patterns)

    # Which patterns to be included                                                                                                        
    parser.add_option("--patterns", dest="patterns_string",
                    help="Limit the result to some CAAS patterns. Patterns are indicated with numbers from 1 to 3 and must \
                        be provided as comma separated (e.g.: -s 1,2,3). See documentation for more details on the patterns", default = "1,2,3")

    ###     1.4.2 filter per gaps 

    # Max background gaps
    parser.add_option("--max_bg_gaps", dest="max_bg_gaps_string",
                    help="Filter by number of gaps in the background\
                        e.g. --max_bg_gaps 3 will accept positions with less than 3 gaps in bg. Default = nofilter (all gaps accepted).", default = "NO")

    # Max foreground gaps
    parser.add_option("--max_fg_gaps", dest="max_fg_gaps_string",
                    help="Filter by number of gaps in the foreground.\
                        e.g. --max_bg_gaps 3 will accept positions with less than 3 gaps in fg. Default = nofilter (all gaps accepted).", default = "NO")

    # Max overall gaps (fg + bg)
    parser.add_option("--max_gaps", dest="max_gaps_string",
                    help="Filter by number of gaps in foreground and background.\
                        e.g. --max_bg_gaps 3 will accept positions with less than 3 gaps in fg and bg. Default = nofilter (all gaps accepted).", default = "NO")

    # Max gaps per position
    parser.add_option("--max_gaps_per_position", dest="max_gaps_pos_string",
                    help="Max gap ratio admitted in a single alignment position.", default = "0.5")

    ###     1.4.3 filter for missing species

    # Max background missing species
    parser.add_option("--max_bg_miss", dest="max_bg_miss_string",
                    help="Filter by number of miss in the background\
                        e.g. --max_bg_miss 3 will accept positions with less than 3 miss in bg. Default = nofilter (all miss accepted).", default = "NO")

    ###     1.4.4 Max foreground missing species
    parser.add_option("--max_fg_miss", dest="max_fg_miss_string",
                    help="Filter by number of miss in the foreground.\
                        e.g. --max_bg_miss 3 will accept positions with less than 3 miss in fg. Default = nofilter (all miss accepted).", default = "NO")

    ###     1.4.5 Max overall missing species (fg + bg)
    parser.add_option("--max_miss", dest="max_miss_string",
                    help="Filter by number of miss in foreground and background.\
                        e.g. --max_bg_miss 3 will accept positions with less than 3 miss in fg and bg. Default = nofilter (all miss accepted).", default = "NO")


    ### 1.5 Usage

    parser.usage = "ct discoevery -a $alignment_file -t $trait_file -o $output_file --fmt $alignment_format (default:clustal)"

    ### 1.6 Parse the options

    (options, args) = parser.parse_args()

    ### 1.7 Check the options

    missing_option_messages = []

    if options.single_alignment == "none":
        missing_option_messages.append("No input MSA file provided")
    
    if options.config_file == "none":
        missing_option_messages.append("No config file provided")

    if options.output_file == "none":
        missing_option_messages.append("Output not specified")
    
    if len(missing_option_messages) > 0:
        print("\n" + application_info)
        print("\n\n****ERROR: mandatory i/o information missing:")
        print("\n".join(missing_option_messages))
        print("")
        print(parser.usage)
        print("")
        print("For further info: ct discovery --help")
        print("")
        exit()


    ### 1.8 Import the modules

    from modules.disco import *
    from modules.runslice import runslice

    ### 1.9 PROCEDURE Step 1- Slice the alignment

    sliced_alignment = runslice(options)

    ### 1.10 PROCEDURE Step 2- Run the discovery

    print(application_info)
    print("")

    print("[DISCOVERY TOOL] - Scanning", options.single_alignment, "with phenotype information from", options.config_file + "\n\n")


    discovery( 
                input_cfg = options.config_file,
                sliced_object = sliced_alignment,

                max_fg_gaps = options.max_fg_gaps_string,
                max_bg_gaps = options.max_bg_gaps_string,
                max_overall_gaps = options.max_gaps_string,

                max_fg_miss = options.max_fg_miss_string,
                max_bg_miss = options.max_bg_miss_string,
                max_overall_miss = options.max_miss_string,

                admitted_patterns = options.patterns_string,
                output_file = options.output_file)

    if exists(options.output_file):
        print("\n\nDone. CAAS discovery table is available at:\n\n\t" + options.output_file + "\n\n")
    else:
        print("\n\nWarning: No CAAS Found, CAAStools generated no output file.\n")




#### TOOL 2. RESAMPLE ######################################################################################################
########################################################################################################################

if tool.lower() == "resample":

    ### 2.1 Check the dependencies
    check_dependencies("resample", ["dendropy"])

    ### 2.2 Init the input parser
    parser = OptionParser()

    ### 2.3 Inputs and outputs

    #       2.3.1 Phylogeny (phylogenetic tree for permulations)
    parser.add_option("-p", "--phylogeny", dest="phylogeny_file",
                    help="The tree of species in newick format. Needed for permulations", default = "none")

    ###     2.3.2 Output file (the table)
    parser.add_option("-o", "--output", dest="output_file",
                    help="The output file with resampled traits", default = "none")

    ### 2.4 FG and BG size settings

    #       2.4.1 Template config file (ct discovery -t)
    parser.add_option("--bytemp", dest="config_file",
                    help="Fetch the size of FG and BG groups from a ct discovery trait config file.", default = "none")

    #       2.4.2 Size of the foreground group
    parser.add_option("-f", "--fg_size", dest="fgsize",
                    help="Number of species in the foreground group. Must be > 0.", default = "notset")

    #       2.4.3 Size of the background group
    parser.add_option("-b", "--bg_size", dest="bgsize",
                    help="Number of species in the background group. Must be > 0.", default = "notset")

    ### 2.5 Simulation settings

    ###     2.5.1 Simulation strategy
    parser.add_option("-m", "--mode", dest="bootstrap_mode",
                    help="Virtualization strategy. 'random' for random species selection, or 'bm' for brownian motion based selection.", default = "random")

    ###     2.5.2 Simulation strategy options for phylogeny restriction
    parser.add_option("--limit_by_group", dest="groupfile",
                    help="Limits random species selections in specific groups. Works with --mode random only and requires a group specification file.", default = "none")

    parser.add_option("--limit_by_patristic_distance", dest="pd_option",
                    help="Limits random species selections in specific patristic distance intervals. Works with --mode random only. Overridden by --limit_by_group.", default = "NO")

    ###     2.5.3 Simulation strategy options for brownian motion

    parser.add_option("--traitvalues", dest="trait_values",
                    help="Trait values for brownian motion reshuffling. Mandatory for --mode bm.", default = "none")

    ###     2.5.4 How many bootstrap cycles?
    parser.add_option("--cycles", dest="cycles",
                    help="number of cycles", default = "1000")

    ### 2.6 Usage

    parser.usage = "ct resample -p $phylogenetic_tree (newick format) -f $foreground_size -b $background_size / --bytemp $trait_file -o $output_file\n\nNOTE: to use --mode bm or phylogeny restriction you MUST provide a template (--bytemp)"

    ### 2.7 Parse the options

    (options, args) = parser.parse_args()

    ### 2.7 Check the mandatory options

    ###     2.7.1 the general input

    missing_option_messages = []

    if options.phylogeny_file == "none":
        missing_option_messages.append("No input tree file provided")
    
    if options.output_file == "none":
        missing_option_messages.append("Output not specified")
    
    if len(missing_option_messages) > 0:
        print("\n" + application_info)
        print("\n\n****ERROR: mandatory i/o information missing:")
        print("\n".join(missing_option_messages))
        print("")
        print(parser.usage)
        print("")
        print("For further info: ct resample --help")
        print("")
        exit()

    ###     2.7.2 The group size determination

    if options.config_file == "none":
        if "notset" in [options.fgsize, options.bgsize ]:
            print("\n\n****ERROR: you need to specify background and foreground size or provide a trait config file as a template")
            print("")
            print(parser.usage)
            print("")
            exit()
        
        else:
            try:
                bgsize_numeric = int(options.bgsize)
                fgsize_numeric = int(options.fgsize)

            except:
                print("\n\n****ERROR: background and/or foreground size in invalid format (must be an integer)" )
                print("")
                print(parser.usage)
                print("")
                exit()

    else:
        with open(options.config_file) as cfg_handle:
            cfg_list = cfg_handle.read().splitlines()
    
        values = []

        for x in cfg_list:
            try:
                c = x.split("\t")
                values.append(c[1])
            except:
                pass
        
        fgsize_numeric = values.count("1")
        bgsize_numeric = values.count("0")
    
    ###     2.7.3 I need a template if you use phylogeny restriction or bm

    if options.bootstrap_mode == "random" and options.pd_option != "NO":
        if options.config_file == "none":
            print("\n\n****ERROR: phylogeny restriction in random mode or brownian motion require a config file as a template.")
            print("")
            print(parser.usage)
            print("")
            exit()

    if options.bootstrap_mode == "random" and options.groupfile != "none":
        if options.config_file == "none":
            print("\n\n****ERROR: phylogeny restriction in random mode or brownian motion require a config file as a template.")
            print("")
            print(parser.usage)
            print("")
            exit()

    if options.bootstrap_mode == "bm" and options.config_file == "none":
            print("\n\n****ERROR: phylogeny restriction in random mode or brownian motion require a config file as a template.")
            print("")
            print(parser.usage)
            print("")
            exit()


    ### 3 Import the bootstrap initialisation

    from modules.init_bootstrap import *

    normalised_mode = options.bootstrap_mode

    if options.bootstrap_mode == "random" and options.groupfile != "none":
        normalised_mode = "phylogeny-restricted-byfams"

    if options.bootstrap_mode == "random" and options.pd_option == "YES":
        normalised_mode = "phylogeny-restricted-bypd"

    print(application_info)
    print("")

    print("[RESAMPLE TOOL] - Simulating traits...")


    w = simtrait(
        fg_len = fgsize_numeric,
        bg_len = bgsize_numeric,
        template = options.config_file,
        tree_file = options.phylogeny_file,
        mode = normalised_mode,
        groupfile = options.groupfile,
        phenotype_values_file = options.trait_values,
        cycles = int(options.cycles),
        simtraits_outfile = options.output_file
    )

    # Output information (recaps the simulation and the settings)

    if normalised_mode == "random":
        print("\n\nTrait simulation in", options.bootstrap_mode, "mode with", options.cycles, "cycles is done. Simulation file is avaiable at:\n\n\t" + options.output_file) 

    elif normalised_mode == "phylogeny-restricted-byfams":
        print("\n\nTrait simulation in", options.bootstrap_mode, "mode, restricted by family information from", options.groupfile, "with", options.cycles, "cycles is done. Simulation file is avaiable at:\n\n\t" + options.output_file)

    elif normalised_mode == "phylogeny-restricted-bypd":
        print("\n\nTrait simulation in", options.bootstrap_mode, "mode, restricted by phylogenetic consistency based on patristic distance intervals", options.groupfile, "with", options.cycles, "cycles is done. Simulation file is avaiable at:\n\n\t" + options.output_file)
    
    elif normalised_mode == "bm":
        if exists(options.output_file):
            print("\n\nTrait simulation in", options.bootstrap_mode, "mode, based on", options.groupfile, "template and ", options.trait_values, "trait values with", options.cycles, "cycles is done. Simulation file is avaiable at:\n\n\t" + options.output_file)
        else:
            print("\n\n****ERROR: resampled traits file not generated. See R interpreter output.\n\n")
    print("\nThis file can be used as input for the bootstrap tool\n\n")




#### TOOL 3. BOOTSTRAP ######################################################################################################
########################################################################################################################

if tool.lower() == "bootstrap":

    ### 3.1 Check the dependencies
    check_dependencies("discovery", ["biopython", "scipy", "numpy"])

    ### 3.2 Init the input parser
    parser = OptionParser()

    ### 3.3 Inputs and outputs

    ###     3.3.1 Config file (phylogenetic tree for permulations)
    parser.add_option("-t", "--traitfile", dest="config_file",
                    help="The trait config file (read documentation for file formatting)", default = "none")

    ###     3.3.2 Resampled traits File (phylogenetic tree for permulations)
    parser.add_option("-s", "--simtraits", dest="simtraits",
                    help="The resampled traits file", default = "none")

    ###     3.3.3 Input alignment
    parser.add_option("-a", "--alignment", dest="single_alignment",
                    help="Multiple Sequence Alignment (MSA) file.", default = "none")

    ###     3.3.4 Alignment format
    parser.add_option("--fmt", dest="ali_format",
                    help="File format of the MSA file. Default: clustal. Accepted: clustal, emboss, fasta, \
                    fasta-m10, ig, maf, mauve, msf, nexus, phylip, phylip-sequential, phylip-relaxed, stockholm.", default = "clustal")    
    ###     3.3.5 Output file (the table)
    parser.add_option("-o", "--output", dest="output_file",
                    help="The output file, where the bootstrap table will be printed", default = "none")
    


    ### 3.4 Optional and filtering inputs

    ###     3.4.1 limit output to pattern (default: all patterns)

    # Which patterns to be included                                                                                                        
    parser.add_option("--patterns", dest="patterns_string",
                    help="Limit the result to some patterns. Patterns are indicated with numbers from 1 to 4 and must \
                        be provided as comma separated (e.g.: -s 1,2,3). See documentation for more details on the patterns", default = "1,2,3")

    ###     3.4.2 filter per gaps 

    # Max background gaps
    parser.add_option("--max_bg_gaps", dest="max_bg_gaps_string",
                    help="Filter by number of gaps in the background\
                        e.g. --max_bg_gaps 3 will accept positions with less than 3 gaps in bg. Default = nofilter (all gaps accepted).", default = "NO")

    # Max foreground gaps
    parser.add_option("--max_fg_gaps", dest="max_fg_gaps_string",
                    help="Filter by number of gaps in the foreground.\
                        e.g. --max_bg_gaps 3 will accept positions with less than 3 gaps in fg. Default = nofilter (all gaps accepted).", default = "NO")

    # Max overall gaps (fg + bg)
    parser.add_option("--max_gaps", dest="max_gaps_string",
                    help="Filter by number of gaps in foreground and background.\
                        e.g. --max_bg_gaps 3 will accept positions with less than 3 gaps in fg and bg. Default = nofilter (all gaps accepted).", default = "NO")

    # Max gaps per position
    parser.add_option("--max_gaps_per_position", dest="max_gaps_pos_string",
                    help="Max gap ratio admitted in a single alignment position.", default = "0.5")

    ###     3.4.3 filter for missing species

    # Max background missing species
    parser.add_option("--max_bg_miss", dest="max_bg_miss_string",
                    help="Filter by number of miss in the background\
                        e.g. --max_bg_miss 3 will accept positions with less than 3 miss in bg. Default = nofilter (all miss accepted).", default = "NO")

    ###     3.4.4 Max foreground missing species
    parser.add_option("--max_fg_miss", dest="max_fg_miss_string",
                    help="Filter by number of miss in the foreground.\
                        e.g. --max_bg_miss 3 will accept positions with less than 3 miss in fg. Default = nofilter (all miss accepted).", default = "NO")

    ###     3.4.5 Max overall missing species (fg + bg)
    parser.add_option("--max_miss", dest="max_miss_string",
                    help="Filter by number of miss in foreground and background.\
                        e.g. --max_bg_miss 3 will accept positions with less than 3 miss in fg and bg. Default = nofilter (all miss accepted).", default = "NO")


    ### 3.5 Usage

    parser.usage = "ct bootstrap -a $alignment_file -t $trait_config_file -s $resampled_traits_file -o $output_file --fmt $alignment_format (default:clustal)"

    ### 3.6 Parse the options

    (options, args) = parser.parse_args()

    ### 3.7 FILTERINGS AND ERRORS

    missing_option_messages = []

    if options.single_alignment == "none":
        missing_option_messages.append("No input MSA file provided")

    if options.config_file == "none":
        missing_option_messages.append("You must provide a trait config file")

    if options.simtraits == "none":
        missing_option_messages.append("No config file provided")

    if options.output_file == "none":
        missing_option_messages.append("Output not specified")
    
    if len(missing_option_messages) > 0:
        print("\n" + application_info)
        print("\n\n****ERROR: mandatory i/o information missing:")
        print("\n".join(missing_option_messages))
        print("")
        print(parser.usage)
        print("")
        print("For further info: ct bootstrap --help")
        print("")
        exit()

    ### 3.8 Import the modules

    from modules.boot import *
    from modules.init_bootstrap import simtrait_revive
    from modules.runslice import runslice

    ### 3.9 PROCEDURE 

    print(application_info)
    print("")

    print("[BOOTSTRAP TOOL] - Scanning", options.single_alignment, "with phenotype information from", options.config_file + "\n\n")

    ###     3.9.1 - Slice the alignment

    sliced_alignment = runslice(options)

    ###     3.9.2 - Read the resampled traits file

    bootstrap_object = simtrait_revive(options.simtraits)

    ###     3.9.3 - Boot on the alignment file
    boot_on_single_alignment(
                    trait_config_file= options.config_file,
                    resampled_traits= bootstrap_object,
                    sliced_object = sliced_alignment,

                    max_fg_gaps = options.max_fg_gaps_string,
                    max_bg_gaps = options.max_bg_gaps_string,
                    max_overall_gaps = options.max_gaps_string,
                    
                    max_fg_miss = options.max_fg_miss_string,
                    max_bg_miss = options.max_bg_miss_string,
                    max_overall_miss = options.max_miss_string,

                    the_admitted_patterns = options.patterns_string,
                    output_file = options.output_file
                    )
    
    ###     3.9.4 Final output
    print("\n\nBootstrap information available in", options.output_file)