#!/opt/anaconda1anaconda2anaconda3/bin/python
# File created on 17 Apr 2013
from __future__ import division

__author__ = "Will Van Treuren"
__copyright__ = "Copyright 2013, The QIIME project"
__credits__ = ["Will Van Treuren, Julia Goodrich, Luke Ursell"]
__license__ = "GPL"
__version__ = "1.8.0"
__maintainer__ = "Will Van Treuren"
__email__ = "wdwvt1@gmail.com"


from qiime.util import parse_command_line_parameters, make_option, create_dir
import os
from qiime.parse import parse_mapping_file_to_dict
from biom.parse import parse_biom_table
from qiime.make_bipartite_network import (make_sample_node_table, _write_table,
    make_otu_node_table, make_node_attr_table, make_edge_table)

script_info = {}
script_info['brief_description'] = """This script makes a bipartite network connecting samples to observations. It is most suitable for visualization with cytoscape."""
script_info['script_description'] = \
"""This script was created to ease the process of making bipartite networks that have appeared in high profile publications including 10.1073/pnas.1217767110 and 10.1126/science.1198719. The script will take an biom table and a mapping file and produce an edge table which connects each sample in the biom table to the observations found in that sample. It is bipartite because there are two distinct node classes -- OTU and Sample. The 'OTU' node class does not have to be an operational taxonomic unit, it can be a KEGG category or metabolite etc. -- anything that is an observation. The edges are weighted by the abundance of the observation in the sample to which it is connected. The output files of this script are intended to be loaded into Cytoscape. The EdgeTable should be uploaded first, and then the NodeAttrTable file can be uploaded as node attributes to control coloring, sizing, and shaping as the user desires. The overall idea behind this script is to make bipartite network creation easier. To that end, the color, size, and shape options are used to provide fields in the NetworkViz tab of Cytoscape so that nodes can be appropriately presented. Those options are passed via comma separated strings (as in the example below). The most common visualization strategy is to color sample nodes by a metadata category like timepoint or pH, color OTU nodes by one of their taxonomic levels, and to scale OTU node size by abundance. This script makes this process easy (as well as a myriad of other visualiation strategies). Once the tables are created by this script they must be opened in cytoscape. This process is described in detail in the QIIME bipartite network tutorial available at: http://qiime.org/tutorials/making_cytoscape_networks.html All color, size, and shape options in this script default to 'NodeType'. OTU nodes have NodeType: otu, sample nodes have NodeType: sample. Thus, if you ran this script with defaults, you would only be able to change the shape, size, and color of the nodes depending on whether or not they were observations or samples. You would not be able to distinguish between two observations based on color, shape, or size. The script is flexible in that it allows you to pass any number of fields for the --{s,o}{shape,size,color}. This will allow you to distinguish between OTU and sample nodes in a huge number of different ways. The usage examples below show some of the common use cases and what options you would pass to emulate them. There are a couple of important considerations for using this script:

Note that the --md_fields option has a different meaning depending on the type of metadata in the biom table. Regardless of type, the md_fields will be the headers in the OTUNodeTable.txt. If the metadata is a dict or default dict, the md_fields will be used as keys to extract data from the biom file metadata. If the metadata is a list or a string, then the md_fields will be have no intrinsic relation to the columns they head. For example if md_fields=['k','p','c'] and the metadata contained in a given OTU was 'k__Bacteria;p__Actinobacter;c__Actino' the resulting OTUNodeTable would have k__Bacteria in the 'k' column, p__Actinobacter in the 'p' column, and c__Actino in the 'c' column. If one passed md_fields=['1.0','XYZ','Five'] then the OTUNodeTable would have columns headed by ['1.0','XYZ','Five'], but the metadata values in those columns would be the same (e.g. '1.0' column entry would be k__Bacteria etc.) If the number of elements in the metadata for a given OTU is not equal to the number of headers provided the script will adjust the OTU metadata. In the case where the metadata is too short, it will add 'Other' into the OTU metadata until the required length is reached.  In the case where the metadata is too long it will simply remove extra entries. This means you can end up with many observations which have the value of 'Other' if you have short taxonomic strings/lists for your observations.

The available fields for both sample and otu nodes are:
[NodeType, Abundance]

For observation nodes the additional fields available are:
any fields you passed for the md_fields

For sample nodes the additional fields available are
any fields found in the mapping file headers

If multiple fields are passed for a given option, they will be concatenated in the output with a '_' character. 
"""
script_info['script_usage'] = [\
("","""Create an EdgeTable and NodeAttrTable that allow you to color sample nodes with one of their metadata categories (Treatment for our example), observation nodes (in this case OTUs) by their taxonomic level (class for our example), control observation node size by their abundance, and control node shape by whether its an observation or sample.""",
"""%prog -i otu_table.biom -m mapping_file.txt -k taxonomy --md_fields 'k,p,c,o,f' -o bipartite_network/ --scolors 'Treatment' --ocolors 'c' --osize 'Abundance'"""),
("","""Create an EdgeTable and NodeAttrTable that allow you to color sample nodes by a combination of their time point and diet, color observation nodes by their abundance and family, and node shape by whether the node is an observation or sample. Note that the names in the --md_fields are irrelevant as long as the field passed for --ocolors is available. The length is important however, since there are 5 levels in our OTU table. If fewer fewer than 5 fields were passed for --md_fields we would get an error.""",
"""%prog -i otu_table.biom -m mapping_file.txt -k taxonomy --md_fields 'a1,a2,a3,a4,f' -o bipartite_network_combo_colors/ --scolors 'TimePt,Diet' --ocolors 'f,Abundance'""")]
script_info['script_usage_output_to_remove'] = ['$PWD/bipartite_network/','$PWD/bipartite_network_combo_colors/']
script_info['output_description']= """The output of this script is four files:
    1. EdgeTable - table with connections between samples and observations. 
    2. OTUNodeTable - table with observations and their associated metadata. 
    3. SampleNodeTable - table with samples and their associated metadata.
    4. NodeAttrTable - table with the node attributes specified by the user with
       the given options."""
script_info['required_options'] = [\
 make_option('-i','--biom_fp',type="existing_filepath",help='the input file ' +\
    'path for biom table.'),
 make_option('-m','--map_fp',type="existing_filepath",help='the input file ' +\
    'path for mapping file.'),
 make_option('-o','--output_dir',type="new_dirpath",help='directory that ' +\
    'will be created for storing the results.'),
 make_option('-k','--observation_md_header_key',type="string",help='Key to retrieve ' +\
    'metadata (usually taxonomy) from the biom file.'),
 make_option('--md_fields',type="string",help='metadata fields that will be '+\
    'the headers of the OTUNodeTable. If the biom table has ' +\
    'metadata dictionaries, md_fields will be the keys extracted from the ' +\
    'biom table metadata. Passed like "kingdom,phylum,class".')]
script_info['optional_options'] = [\
 make_option('--scolors',type="string",help='commas separated string ' +\
    'specifying fields of interest for sample node coloring [default: %default].',
    default='NodeType'),
 make_option('--ocolors',type="string",help='commas separated string ' +\
    'specifying fields of interest for observation node coloring [default: %default].',
    default='NodeType'),
 make_option('--sshapes',type="string",help='commas separated string ' +\
    'specifying fields of interest for sample node shape [default: %default].',
    default='NodeType'),
 make_option('--oshapes',type="string",help='commas separated string ' +\
    'specifying fields of interest for observation node shape [default: %default].',
    default='NodeType'),
 make_option('--ssizes',type="string",help='commas separated string ' +\
    'specifying fields of interest for sample node size [default: %default].',
    default='NodeType'),
 make_option('--osizes',type="string",help='commas separated string ' +\
    'specifying fields of interest for observation node size [default: %default].',
    default='NodeType')]


script_info['version'] = __version__

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

    otu_table_fp = opts.biom_fp
    map_fp = opts.map_fp
    output_dir = opts.output_dir
    taxonomy_key = opts.observation_md_header_key
    scolors = opts.scolors.split(',')
    ocolors = opts.ocolors.split(',')
    sshapes = opts.sshapes.split(',')
    oshapes = opts.oshapes.split(',')
    ssizes = opts.ssizes.split(',')
    osizes = opts.osizes.split(',')
    md_fields = opts.md_fields.split(',')

    # check that the otu fields asked for are available
    shared_options = ['NodeType','Abundance']
    if not all([i in md_fields+shared_options for i in ocolors+oshapes+osizes]):
        option_parser.error('The fields specified for observation colors, sizes, or '+\
            'shapes are not in either the shared options (NodeType,Abundance)'+\
            ' or the supplied md_fields. These fields must be a subset of '+\
            'the union of these sets. Have you passed ocolors, osizes or '+\
            'oshapes that are not in the md_fields?')
    # check that the sample fields asked for are available. mapping file 
    # elements should all have same metadata keys
    sopts = parse_mapping_file_to_dict(map_fp)[0].items()[0][1].keys()
    if not all([i in sopts+shared_options for i in scolors+sshapes+ssizes]):
        option_parser.error('The fields specified for sample colors, sizes, or '+\
            'shapes are not in either the shared options (NodeType,Abundance)'+\
            ' or the supplied mapping file. These fields must be a subset of '+\
            'the union of these sets. Have you passed scolors, ssizes or '+\
            'sshapes that are not in the mapping file headers?')

    # actual compuation begins
    try:
        create_dir(output_dir, fail_on_exist=True)
    except OSError:
        option_parser.error('Directory already exists. Will not overwrite.')

    bt = parse_biom_table(open(otu_table_fp))
    pmf = parse_mapping_file_to_dict(map_fp)[0] # [1] is comments, don't need
    sample_node_table = make_sample_node_table(bt, pmf)
    otu_node_table = make_otu_node_table(bt, opts.observation_md_header_key, 
        md_fields)
    node_attr_table = make_node_attr_table(otu_node_table, sample_node_table,
        scolors, ocolors, ssizes, osizes, sshapes, oshapes)
    edge_table = make_edge_table(bt)

    _write_table(sample_node_table, os.path.join(output_dir,'SampleNodeTable.txt'))
    _write_table(otu_node_table, os.path.join(output_dir,'OTUNodeTable.txt'))
    _write_table(node_attr_table, os.path.join(output_dir,'NodeAttrTable.txt'))
    _write_table(edge_table, os.path.join(output_dir,'EdgeTable.txt'))

if __name__ == "__main__":
    main()