#!/usr/bin/python
#-*- coding: UTF-8 -*-

"""
 (c) 2011, 2012 - Copyright Pierre-Yves Chibon

 Distributed under License GPLv3 or later
 You can find a copy of this license on the website
 http://www.gnu.org/licenses/gpl.html

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 3 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 MA 02110-1301, USA.
"""


import argparse
import logging
import os
import shutil
import sys
import tempfile

try:
    from MQ2 import (set_tmp_folder, extract_zip, MQ2Exception,
        MQ2NoMatrixException)
    from MQ2.generate_map_from_mapqtl import generate_map_from_mapqtl
    from MQ2.parse_mapqtl_file import parse_mapqtl_file
    from MQ2.add_marker_to_qtls import add_marker_to_qtls
    from MQ2.add_qtl_to_map import add_qtl_to_map
except:
    from src import (set_tmp_folder, extract_zip, MQ2Exception,
        MQ2NoMatrixException)
    from src.generate_map_from_mapqtl import generate_map_from_mapqtl
    from src.parse_mapqtl_file import parse_mapqtl_file
    from src.add_marker_to_qtls import add_marker_to_qtls
    from src.add_qtl_to_map import add_qtl_to_map

# Initial simple logging stuff
logging.basicConfig()
LOG = logging.getLogger('MQ2')

if '--debug' in sys.argv:
    LOG.setLevel(logging.DEBUG)
elif '--verbose' in sys.argv:
    LOG.setLevel(logging.INFO)


def get_arguments():
    """ Handle the command line arguments given to this program """
    LOG.debug('Parse command line argument')
    parser = argparse.ArgumentParser(description='Command line interface \
        for the MQ² program')
    parser.add_argument('-z', '--zipfile', default=None,
        help='Zip file containing the MapQTL output files.')
    parser.add_argument('-f', '--folder', default=None,
        help='Path to a local folder containing the MapQTL output files.')
    parser.add_argument('--lod', default=3,
        help='LOD threshold to use to assess the significance of a LOD \
        value for a QTL.')
    parser.add_argument('--session', default=None,
        help='MapQTL session to analyze.')
    parser.add_argument('--verbose', action='store_true',
                help="Gives more info about what's going on")
    parser.add_argument('--debug', action='store_true',
                help="Outputs debugging information")
    return parser.parse_args()


def main():
    """ Main function.
    Retrieves the arguments provided by the user.
    Extract the Zip file if needed.
    Convert the .map file to CSV.
    Parse the output files from MapQTL
        Generate the qtl_matrix.csv
        Generate the qtls.csv file
    Retrieve the marker closest to the QTL peak
        Generate the qtl_with_mk.csv file
    Calculate for each markers the number of QTLs associated
        Generate the map_with_qtl.csv file
    """
    args = get_arguments()
    if args.zipfile and args.folder:
        print 'You cannot specify and a zip file and a folder for the '\
            'MapQTL output.'
        return 1
    if not args.zipfile and not args.folder:
        print 'You must provide either a zip file or a folder '\
            'containing the MapQTL output.'
        return 1
    if not args.session:
        print 'You must provide the MapQTL session to analyse.'
        return 1

    tmp_folder = None
    try:
        if args.zipfile:
            mapqtl_folder = set_tmp_folder()
            extract_zip(args.zipfile, mapqtl_folder)
        else:
            mapqtl_folder = args.folder

        generate_map_from_mapqtl(inputfolder=mapqtl_folder,
                sessionid=args.session)

        try:
            parse_mapqtl_file(inputfolder=mapqtl_folder,
                sessionid=args.session,
                lodthreshold=args.lod,
                qtl_outputfile='qtls.csv')
        except MQ2NoMatrixException, err:
            LOG.warn(err)

        add_marker_to_qtls(qtlfile='qtls.csv',
            mapfile='map.csv',
            outputfile='qtls_with_mk.csv')

        add_qtl_to_map(qtlfile='qtls_with_mk.csv',
            mapfile='map.csv',
            outputfile='map_with_qtl.csv')
    except MQ2Exception, err:
        print 'ERROR: %s' % err
        return 2
    finally:
        if tmp_folder and os.path.exists(tmp_folder):
            shutil.rmtree(tmp_folder)


if __name__ == "__main__":
    main()
