#!/opt/anaconda1anaconda2anaconda3/bin/python
# File created on 02 Nov 2012
from __future__ import division

__author__ = "Yoshiki Vazquez-Baeza"
__copyright__ = "Copyright 2011, The QIIME project"
__credits__ = ["Yoshiki Vazquez-Baeza", "Antonio Gonzalez-Pena"]
__license__ = "GPL"
__version__ = "1.8.0"
__maintainer__ = "Yoshiki Vazquez-Baeza"
__email__ = "yoshiki89@gmail.com"

from os.path import basename, splitext
from qiime.format import format_mapping_file
from qiime.parse import parse_matrix, parse_mapping_file
from qiime.util import parse_command_line_parameters, make_option
from qiime.add_alpha_to_mapping_file import (mean_alpha,
    add_alpha_diversity_values_to_mapping_file)

script_info = {}
script_info['brief_description'] = "Add alpha diversity data to a metadata "+\
    "mapping file"
script_info['script_description'] = "Add alpha diversity data to a mapping "+\
    "file for use with other QIIME scripts, i. e. make_emperor.py. The "+\
    "resulting mapping file will contain three new columns per metric in the "+\
    "alpha diversity data; the first column being the raw value, the second "+\
    "being a normalized raw value and the third one a label classifying "+\
    "the bin where this value fits based on the normalized value."
script_info['script_usage'] = [("Adding alpha diversity data:","Add the alpha"
    " diversity values to a mapping file and classify the normalized values "
    "into 4 bins, where the limits will be  0 < x <= 0.25 for the first bin "
    "0.25 < x <= 0.5 for the second bin, 0.5 < x <= 0.75 for the third bin "
    "and 0.75 < x <= 1 for the fourth bin.","%prog -i adiv_pd.txt -m mapping"
    ".txt -b 4 -o alpha_mapping.txt")]
script_info['script_usage'].append(("Adding alpha diversity data with the "
    "quantile method:", "Add the alpha diversity values to a mapping file and"
    " classify the normalized values using the quartiles of the distribution "
    "of this values.", "%prog -i adiv_pd.txt -m mapping.txt -b 4 -o alpha_"
    "mapping_quantile.txt --binning_method=quantile"))
script_info['script_usage'].append(("Adding collated alpha diversity data",
    "Add the mean of the alpha diversity values at a specified rarefaction"
    " depth, this case is for use with the output of collate_alpha.py. It is "
    "recommended that the filenames are the name of the metric used in each "
    "file.", "%prog -i 'shannon.txt,chao1.txt' -m mapping.txt -b 4 -o collated_"
    "alpha_mapping.txt --depth=49 --collated_input"))
script_info['output_description']= "The result of running this script is a "+\
    "metadata mapping file that will include 3 new columns per alpha "+\
    "diversity metric included in the alpha diversity file. For example, with"+\
    " an alpha diversity file with only PD_whole_tree, the new columns will "+\
    "PD_whole_tree_alpha, PD_whole_tree_normalized and PD_whole_tree_bin."
script_info['required_options'] = [
make_option('-i','--alpha_fps',type="existing_filepaths",
    help='alpha diversity data with one or multiple metrics i. e. the output'
    ' of alpha_diversity.py. This can also be a comma-separated list of colla'
    'ted alpha diversity file paths i. e. the output of collate_alpha.py, when'
    ' using collated alpha diversity data the --depth option is required'),
make_option('-m','--mapping_fp',type="existing_filepath",
    help='mapping file to modify by adding the alpha diversity data'),
]
script_info['optional_options'] = [
make_option('-o','--output_mapping_fp',type="new_filepath",
    help='filepath for the modified mapping file [default: %default]',
    default='mapping_file_with_alpha.txt'),
make_option('-b','--number_of_bins',type="int",
    help='number of bins [default: %default].', default=4),
make_option('-x','--missing_value_name',type="string",
    help='bin prefix name for the sample identifiers that exist in the '
        'mapping file (mapping_fp) but not in the alpha diversity file '
        '(alpha_fp) [default: %default].', default='N/A'),
make_option('--binning_method', type='choice', choices=['equal', 'quantile'],
    default='equal',
    help='Select the method name to create the bins, the options are '
    '\'equal\' and \'quantile\'. Both methods work over the normalized alpha '
    'diversity values. On the one hand \'equal\' will assign the bins on '
    'equally spaced limits, depending on the value of --number_of_bins i. e. '
    'if you select 4 the limits will be [0.25, 0.50, 0.75]. On the other hand '
    '\'quantile\' will select the limits based on the --number_of_bins i. e. '
    'the limits will be the quartiles if 4 is selected [default: %default].'),
make_option('--depth', type='int', default=None, help='Select the rarefaction '
    'depth to use when the alpha_fps refers to collated alpha diversity file(s)'
    ' i. e. the output of collate_alpha.py. All the iterations contained at '
    'this depth will be averaged to form a single mean value [default: highest '
    'depth available].'),
make_option('--collated_input', action='store_true', help='Use to specify that '
    'the -i option is composed of collated alpha diversity data.',
    default=False)
]
script_info['version'] = __version__



def main():
    option_parser, opts, args = parse_command_line_parameters(**script_info)

    alpha_fps = opts.alpha_fps
    mapping_fp = opts.mapping_fp
    output_mapping_fp = opts.output_mapping_fp
    binning_method = opts.binning_method
    missing_value_name = opts.missing_value_name
    depth = opts.depth
    number_of_bins = opts.number_of_bins
    collated_input = opts.collated_input

    # if using collated data, make sure they specify a depth
    if collated_input:
        alpha_dict = {}

        # build up a dictionary with the filenames as keys and lines as values
        for single_alpha_fp in alpha_fps:
            alpha_dict[splitext(basename(single_alpha_fp))[0]] = open(
                single_alpha_fp, 'U').readlines()

        # format the collated data
        try:
            metrics, alpha_sample_ids, alpha_data = mean_alpha(alpha_dict,
                depth)
        except ValueError, e: # see mean_alpha for the possible exceptions
            option_parser.error(e.message)

    # when not using collated data, the user can only specify one input file
    else:
        if len(alpha_fps) > 1:
            option_parser.error('A comma-separated list of files should only be'
                ' passed with the --alpha_fps option when using collated alpha '
                'diversity data and also selecting a rarefaction depth with the'
                ' --depth option.')
        else:
            metrics, alpha_sample_ids, alpha_data = parse_matrix(open(
                alpha_fps[0], 'U'))

    # parse the data from the files
    mapping_file_data, mapping_file_headers, comments = parse_mapping_file(
        open(mapping_fp, 'U'))

    # add the alpha diversity data to the mapping file
    out_mapping_file_data, out_mapping_file_headers = \
        add_alpha_diversity_values_to_mapping_file(metrics, alpha_sample_ids,
        alpha_data, mapping_file_headers, mapping_file_data, number_of_bins,
        binning_method, missing_value_name)

    # format the new data and write it down
    lines = format_mapping_file(out_mapping_file_headers, out_mapping_file_data)
    fd_out = open(output_mapping_fp, 'w')
    fd_out.writelines(lines)
    fd_out.close()

if __name__ == "__main__":
    main()