#!python
#-*- coding: utf-8 -*-
"""
TEA command line script.

Created on Sun Mar 13 16:33:49 2016
@author: dangeles at caltech edu
"""

import os
import pandas as pd
import sys
import tissue_enrichment_analysis.hypergeometricTests as hgt
import argparse
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import re
import seaborn as sns
sns.set_context('paper')
sns.set_style('whitegrid')

if __name__ == '__main__':

    path = './'
    os.chdir(path)

    defQ = 0.1

    parser = argparse.ArgumentParser(description='Run EA.')
    parser = argparse.ArgumentParser()
    parser.add_argument("gene_list",
                        help='The full path to the gene list (WBIDs) you would\
                         like to analyse in .csv format')
    parser.add_argument('title', help='Title for your analysis (shouldn\'t\
                        include file extension)',
                        type=str)
    parser.add_argument('kind', help='What kind of analysis will be ' +
                        'performed. One of `tissue`, `phenotype` or `go`',
                        type=str)
    parser.add_argument("-d", '--dictionary', nargs='?', help='Provide a\
                        dictionary to test. If none given, WormBase URL \
                        will be used to download the corresponding file')
    parser.add_argument("-q", help='Qvalue threshold for significance. \
                        Default is {0} if not provided'.format(defQ),
                        type=float)
    parser.add_argument('-p', '--print', help='Indicate whether you would like \
                        to print results', action='store_true')
    parser.add_argument('-s', "--save", help='Indicate whether to save your \
                        plot.', action='store_true')
    parser.add_argument('-b', "--background", help='Provide a background gene \
                        set as a csv file with a single column without a \
                        column name. Gene names must be in wbid format.',
                        type=str)
    parser.add_argument('-m', "--melted_name", help='Name for gene_to_terms \
                        file. If none provided, defaults to gene_to_terms.csv',
                        type=str)
    args = parser.parse_args()

    gl_name = args.gene_list
    title = args.title

    # optional args
    # load dictionary:
    if args.dictionary:
        dict_name = args.tissue_dictionary
        dictionary = pd.read_csv(dict_name)
    else:
        dictionary = hgt.fetch_dictionary(analysis=args.kind)

    # reduce wbids to background set:
    if args.background:
        bg = pd.read_csv(args.background, header=None, names=['wbid'])
        dictionary = dictionary[dictionary.wbid.isin(bg.wbid)]

    # warn user if the dictionary is empty after subsetting:
    if len(dictionary) == 0:
        raise ValueError('Dictionary is empty after subsetting')

    # set threshold
    if args.q:
        q = args.q
    else:
        q = defQ

    # print results
    if args.print:
        prnt = True
    else:
        prnt = False

    # save results
    if args.save:
        save = True

    else:
        save = False

    # open gene list:
    gene_list = pd.read_csv(gl_name, header=None, names=['wbid'])

    # perform enrichment analysis:
    df_results = hgt.enrichment_analysis(gene_list.wbid.unique(), dictionary,
                                         alpha=q, show=False)

    dfname = title + '.tsv'
    df_results.to_csv(dfname, index=False, sep='\t')

    # melt dictionary:
    melted_dict = pd.melt(dictionary, id_vars='wbid', var_name='term',
                          value_name='found')
    melted_dict = melted_dict[melted_dict.found == 1]

    # keep only terms that were significant:
    melted_dict = melted_dict[melted_dict.term.isin(df_results.Term.values)]

    # keep only relevant genes:
    melted_dict = melted_dict[melted_dict.wbid.isin(gene_list.wbid)]

    # save to file:
    if args.melted_name:
        melted_dict.to_csv(args.melted_name, index=False)
    else:
        melted_dict.to_csv('gene_to_terms.csv', index=False)

    if prnt:
        with open(dfname, 'r') as f:
            printer = f.readlines()

        for value in printer:
            value = value.split(',')
            for val in value:
                if re.findall("\d+\.\d+", val):
                    ind = value.index(val)
                    x = float(val)
                    value[ind] = '{0:.2g}'.format(x)

            value = '\t'.join(value)
            print(value)

    if save:
        hgt.plot_enrichment_results(df_results, title=title, save=save,
                                    analysis=args.kind)

    sys.exit()
