#! python
# -*- coding: utf-8 -*-

import os, sys, logging, json
from collections import OrderedDict
from multiprocessing import Pool, cpu_count
from tempfile import TemporaryDirectory
from numpy import log, exp
import numpy as np
import pandas as pd
import requests
from gseapy.algorithm import enrichment_score, gsea_compute, ranking_metric
from gseapy.algorithm import enrichment_score_tensor, gsea_compute_tensor
from gseapy.parser import gsea_edb_parser, gsea_cls_parser
from gseapy.plot import gseaplot, heatmap
from gseapy.utils import mkdirs, log_init, retry, DEFAULT_LIBRARY, DEFAULT_CACHE_PATH

class GSEAbase(object):
    """base class of GSEA."""
    def __init__(self):
        self.outdir='temp_gseapy'
        self.gene_sets='KEGG_2016'
        self.fdr=0.05
        self.module='base'
        self.results=None
        self.res2d=None
        self.ranking=None
        self.ascending=False
        self.verbose=False
        self._processes=1
        self._logger=None

    def prepare_outdir(self):
        """create temp directory."""
        self._outdir = self.outdir
        if self._outdir is None:
            self._tmpdir = TemporaryDirectory()
            self.outdir = self._tmpdir.name
        elif isinstance(self.outdir, str):
            mkdirs(self.outdir)
        else:
            raise Exception("Error parsing outdir: %s"%type(self.outdir))

        # handle gmt type
        if isinstance(self.gene_sets, str):
            _gset = os.path.split(self.gene_sets)[-1].lower().rstrip(".gmt")
        elif isinstance(self.gene_sets, dict):
            _gset = "blank_name"
        else:
            raise Exception("Error parsing gene_sets parameter for gene sets")

        logfile = os.path.join(self.outdir, "gseapy.%s.%s.log" % (self.module, _gset))
        return logfile

    def _set_cores(self):
        """set cpu numbers to be used"""

        cpu_num = cpu_count()-1
        if self._processes > cpu_num:
            cores = cpu_num
        elif self._processes < 1:
            cores = 1
        else:
            cores = self._processes
        # have to be int if user input is float
        self._processes = int(cores)

    def _load_ranking(self, rnk):
        """Parse ranking file. This file contains ranking correlation vector( or expression values)
           and gene names or ids.

            :param rnk: the .rnk file of GSEA input or a Pandas DataFrame, Series instance.
            :return: a Pandas Series with gene name indexed rankings

        """
        # load data
        if isinstance(rnk, pd.DataFrame):
            rank_metric = rnk.copy()
            # handle dataframe with gene_name as index.
            if rnk.shape[1] == 1: rank_metric = rnk.reset_index()
        elif isinstance(rnk, pd.Series):
            rank_metric = rnk.reset_index()
        elif os.path.isfile(rnk):
            rank_metric = pd.read_csv(rnk, header=None, comment='#', sep="\t")
        else:
            raise Exception('Error parsing gene ranking values!')
        # sort ranking values from high to low
        rank_metric.sort_values(by=rank_metric.columns[1], ascending=self.ascending, inplace=True)
        # drop na values
        if rank_metric.isnull().any(axis=1).sum() >0:
            self._logger.warning("Input gene rankings contains NA values(gene name and ranking value), drop them all!")
            # print out NAs
            NAs = rank_metric[rank_metric.isnull().any(axis=1)]
            self._logger.debug('NAs list:\n'+NAs.to_string())
            rank_metric.dropna(how='any', inplace=True)
        # drop duplicate IDs, keep the first
        if rank_metric.duplicated(subset=rank_metric.columns[0]).sum() >0:
            self._logger.warning("Input gene rankings contains duplicated IDs, Only use the duplicated ID with highest value!")
            # print out duplicated IDs.
            dups = rank_metric[rank_metric.duplicated(subset=rank_metric.columns[0])]
            self._logger.debug('Dups list:\n'+dups.to_string())
            rank_metric.drop_duplicates(subset=rank_metric.columns[0], inplace=True, keep='first')
        # reset ranking index, because you have sort values and drop duplicates.
        rank_metric.reset_index(drop=True, inplace=True)
        rank_metric.columns = ['gene_name','rank']
        rankser = rank_metric.set_index('gene_name')['rank']
        self.ranking = rankser
        # return series
        return rankser

    def load_gmt(self, gene_list, gmt):
        """load gene set dict"""

        if isinstance(gmt, dict):
            genesets_dict = gmt
        elif isinstance(gmt, str):
            genesets_dict = self.parse_gmt(gmt)
        else:
            raise Exception("Error parsing gmt parameter for gene sets")

        subsets = list(genesets_dict.keys())
        self.n_genesets = len(subsets)
        for subset in subsets:
            subset_list = genesets_dict.get(subset)
            if isinstance(subset_list, set):
                subset_list = list(subset_list)
                genesets_dict[subset] = subset_list
            tag_indicator = np.in1d(gene_list, subset_list, assume_unique=True)
            tag_len = tag_indicator.sum()
            if  self.min_size <= tag_len <= self.max_size: continue
            del genesets_dict[subset]

        filsets_num = len(subsets) - len(genesets_dict)
        self._logger.info("%04d gene_sets have been filtered out when max_size=%s and min_size=%s"%(filsets_num, self.max_size, self.min_size))

        if filsets_num == len(subsets):
            self._logger.error("No gene sets passed through filtering condition!!!, try new parameters again!\n" +\
                               "Note: check gene name, gmt file format, or filtering size." )
            sys.exit(0)

        self._gmtdct=genesets_dict
        return genesets_dict

    def parse_gmt(self, gmt):
        """gmt parser"""

        if gmt.lower().endswith(".gmt"):
            with open(gmt) as genesets:
                 genesets_dict = { line.strip().split("\t")[0]: line.strip().split("\t")[2:]
                                  for line in genesets.readlines()}
            return genesets_dict

        elif gmt in DEFAULT_LIBRARY:
            pass
        elif gmt in self.get_libraries():
            pass
        else:
            self._logger.error("No supported gene_sets: %s"%gmt)
            sys.exit(0)

        tmpname = "enrichr." + gmt + ".gmt"
        tempath = os.path.join(DEFAULT_CACHE_PATH, tmpname)
        # if file already download
        if os.path.isfile(tempath):
            self._logger.info("Enrichr library gene sets already downloaded in: %s, use local file"%DEFAULT_CACHE_PATH)
            return self.parse_gmt(tempath)
        else:
            return self._download_libraries(gmt)

    def get_libraries(self, database=''):
        """return active enrichr library name.Offical API """

        lib_url='http://amp.pharm.mssm.edu/%sEnrichr/datasetStatistics'%database
        libs_json = json.loads(requests.get(lib_url).text)
        libs = [lib['libraryName'] for lib in libs_json['statistics']]
        return sorted(libs)

    def _download_libraries(self, libname):
        """ download enrichr libraries."""
        self._logger.info("Downloading and generating Enrichr library gene sets......")
        s = retry(5)
        # queery string
        ENRICHR_URL = 'http://amp.pharm.mssm.edu/Enrichr/geneSetLibrary'
        query_string = '?mode=text&libraryName=%s'
        # get
        response = s.get( ENRICHR_URL + query_string % libname, timeout=None)
        if not response.ok:
            raise Exception('Error fetching enrichment results, check internet connection first.')
        # reformat to dict and save to disk
        mkdirs(DEFAULT_CACHE_PATH)
        genesets_dict = {}
        outname = "enrichr.%s.gmt"%libname
        gmtout = open(os.path.join(DEFAULT_CACHE_PATH, outname), "w")
        for line in response.iter_lines(chunk_size=1024, decode_unicode='utf-8'):
            line=line.strip()
            k = line.split("\t")[0]
            v = list(map(lambda x: x.split(",")[0], line.split("\t")[2:]))
            genesets_dict.update({ k: v})
            outline = "%s\t\t%s\n"%(k, "\t".join(v))
            gmtout.write(outline)
        gmtout.close()

        return genesets_dict

    def _heatmat(self, df, classes, pheno_pos, pheno_neg):
        """only use for gsea heatmap"""
        width = len(classes) if len(classes) >= 6 else  5
        cls_booA =list(map(lambda x: True if x == pheno_pos else False, classes))
        cls_booB =list(map(lambda x: True if x == pheno_neg else False, classes))
        datA = df.loc[:, cls_booA]
        datB = df.loc[:, cls_booB]
        datAB=pd.concat([datA,datB], axis=1)
        self._width = width
        self.heatmat = datAB
        return

    def _plotting(self, rank_metric, results, graph_num, outdir,
                  format, figsize, pheno_pos='', pheno_neg=''):
        """ Plotting API.
            :param rank_metric: sorted pd.Series with rankings values.
            :param results: self.results
            :param data: preprocessed expression table

        """

        # no values need to be returned
        if self._outdir is None: return
        #Plotting
        top_term = self.res2d.index[:graph_num]
        # multi-threading
        pool = Pool(self._processes)
        for gs in top_term:
            hit = results.get(gs)['hits_indices']
            NES = 'nes' if self.module != 'ssgsea' else 'es'
            term = gs.replace('/','_').replace(":","_")
            outfile = '{0}/{1}.{2}.{3}'.format(self.outdir, term, self.module, self.format)
            # gseaplot(rank_metric=rank_metric, term=term, hits_indices=hit,
            #           nes=results.get(gs)[NES], pval=results.get(gs)['pval'],
            #           fdr=results.get(gs)['fdr'], RES=results.get(gs)['RES'],
            #           pheno_pos=pheno_pos, pheno_neg=pheno_neg, figsize=figsize,
            #           ofname=outfile)
            pool.apply_async(gseaplot, args=(rank_metric, term, hit, results.get(gs)[NES],
                                              results.get(gs)['pval'],results.get(gs)['fdr'],
                                              results.get(gs)['RES'],
                                              pheno_pos, pheno_neg,
                                              figsize, 'seismic', outfile))
            if self.module == 'gsea':
                outfile2 = "{0}/{1}.heatmap.{2}".format(self.outdir, term, self.format)
                # heatmap(df=self.heatmat.iloc[hit, :], title=term, ofname=outfile2,
                #         z_score=0, figsize=(self._width, len(hit)/2))
                pool.apply_async(heatmap, args=(self.heatmat.iloc[hit, :], 0, term,
                                               (self._width, len(hit)/2+2), 'RdBu_r',
                                                True, True, outfile2))
        pool.close()
        pool.join()

    def _save_results(self, zipdata, outdir, module, gmt, rank_metric, permutation_type):
        """reformat gsea results, and save to txt"""

        res = OrderedDict()
        for gs, gseale, ind, RES in zipdata:
            rdict = OrderedDict()
            rdict['es'] = gseale[0]
            rdict['nes'] = gseale[1]
            rdict['pval'] = gseale[2]
            rdict['fdr'] = gseale[3]
            rdict['geneset_size'] = len(gmt[gs])
            rdict['matched_size'] = len(ind)
            #reformat gene list.
            _genes = rank_metric.index.values[ind]
            rdict['genes'] = ";".join([ str(g).strip() for g in _genes ])

            if self.module != 'ssgsea':
                # extract leading edge genes
                if rdict['es'] > 0:
                    # RES -> ndarray, ind -> list
                    idx = RES.argmax()
                    ldg_pos = list(filter(lambda x: x<= idx, ind))
                elif rdict['es'] < 0:
                    idx = RES.argmin()
                    ldg_pos = list(filter(lambda x: x >= idx, ind))
                else:
                    ldg_pos = ind # es == 0 ?
                rdict['ledge_genes'] = ';'.join(list(map(str,rank_metric.iloc[ldg_pos].index)))

            rdict['RES'] = RES
            rdict['hits_indices'] = ind
            # save to one odict
            res[gs] = rdict
        # save
        self.results  = res
        # save to dataframe
        res_df = pd.DataFrame.from_dict(res, orient='index')
        res_df.index.name = 'Term'
        res_df.drop(['RES','hits_indices'], axis=1, inplace=True)
        res_df.sort_values(by=['fdr','pval'], inplace=True)
        self.res2d = res_df

        if self._outdir is None: return
        out = os.path.join(outdir,'gseapy.{b}.{c}.report.csv'.format(b=module, c=permutation_type))
        if self.module == 'ssgsea':
            out = out.replace(".csv",".txt")
            with open(out, 'a') as f:
                f.write('# normalize enrichment scores by random permutation procedure (GSEA method)\n')
                f.write("# might not proper for publication\n")
                res_df.to_csv(f, sep='\t')
        else:
            res_df.to_csv(out)

        return

class GSEA(GSEAbase):
    """GSEA main tool"""
    def __init__(self, data, gene_sets, classes, outdir='GSEA_ouput',
                 min_size=15, max_size=500, permutation_num=1000,
                 weighted_score_type=1, permutation_type='gene_set',
                 method='log2_ratio_of_classes', ascending=False,
                 processes=1, figsize=(6.5,6), format='pdf', graph_num=20,
                 no_plot=False, seed=None, verbose=False):

        self.data = data
        self.gene_sets=gene_sets
        self.classes=classes
        self.outdir=outdir
        self.permutation_type=permutation_type
        self.method=method
        self.min_size=min_size
        self.max_size=max_size
        self.permutation_num=int(permutation_num) if int(permutation_num) > 0 else 0
        self.weighted_score_type=weighted_score_type
        self.ascending=ascending
        self._processes=processes
        self.figsize=figsize
        self.format=format
        self.graph_num=int(graph_num)
        self.seed=seed
        self.verbose=bool(verbose)
        self.module='gsea'
        self.ranking=None
        self._noplot=no_plot
        # init logger
        logfile = self.prepare_outdir()
        self._logger = log_init(outlog=logfile,
                                log_level=logging.INFO if self.verbose else logging.WARNING)

    def load_data(self, cls_vec):
        """pre-processed the data frame.new filtering methods will be implement here.
        """
        # read data in
        if isinstance(self.data, pd.DataFrame) :
            exprs = self.data.copy()
            # handle index is gene_names
            if exprs.index.dtype == 'O':
                exprs = exprs.reset_index()
        elif os.path.isfile(self.data) :
            # GCT input format?
            if self.data.endswith("gct"):
                exprs = pd.read_csv(self.data, skiprows=1, comment='#',sep="\t")
            else:
                exprs = pd.read_csv(self.data, comment='#',sep="\t")
        else:
            raise Exception('Error parsing gene expression DataFrame!')

        #drop duplicated gene names
        if exprs.iloc[:,0].duplicated().sum() > 0:
            self._logger.warning("Warning: dropping duplicated gene names, only keep the first values")
            exprs.drop_duplicates(subset=exprs.columns[0], inplace=True) #drop duplicate gene_names.
        if exprs.isnull().any().sum() > 0:
            self._logger.warning("Warning: Input data contains NA, filled NA with 0")
            exprs.dropna(how='all', inplace=True) #drop rows with all NAs
            exprs = exprs.fillna(0)
        # set gene name as index
        exprs.set_index(keys=exprs.columns[0], inplace=True)
        # select numberic columns
        df = exprs.select_dtypes(include=[np.number])
        # drop any genes which std ==0
        df_std =  df.groupby(by=cls_vec, axis=1).std()
        df =  df[~df_std.isin([0]).any(axis=1)]
        df = df + 0.00001 # we don't like zeros!!!
        return df

    def run(self):
        """GSEA main procedure"""

        assert self.permutation_type in ["phenotype", "gene_set"]
        assert self.min_size <= self.max_size

        # Start Analysis
        self._logger.info("Parsing data files for GSEA.............................")
        # phenotype labels parsing
        phenoPos, phenoNeg, cls_vector = gsea_cls_parser(self.classes)
        # select correct expression genes and values.
        dat = self.load_data(cls_vector)
        # data frame must have length > 1
        assert len(dat) > 1
        # ranking metrics calculation.
        dat2 = ranking_metric(df=dat, method=self.method, pos=phenoPos, neg=phenoNeg,
                              classes=cls_vector, ascending=self.ascending)
        self.ranking = dat2
        # filtering out gene sets and build gene sets dictionary
        gmt = self.load_gmt(gene_list=dat2.index.values, gmt=self.gene_sets)

        self._logger.info("%04d gene_sets used for further statistical testing....."% len(gmt))
        self._logger.info("Start to run GSEA...Might take a while..................")
        # cpu numbers
        self._set_cores()
        # compute ES, NES, pval, FDR, RES
        dataset = dat if self.permutation_type =='phenotype' else dat2
        gsea_results,hit_ind,rank_ES, subsets = gsea_compute_tensor(data=dataset, gmt=gmt, n=self.permutation_num,
                                                             weighted_score_type=self.weighted_score_type,
                                                             permutation_type=self.permutation_type,
                                                             method=self.method,
                                                             pheno_pos=phenoPos, pheno_neg=phenoNeg,
                                                             classes=cls_vector, ascending=self.ascending,
                                                             processes=self._processes, seed=self.seed)

        self._logger.info("Start to generate GSEApy reports and figures............")
        res_zip = zip(subsets, list(gsea_results), hit_ind, rank_ES)
        self._save_results(zipdata=res_zip, outdir=self.outdir, module=self.module,
                                   gmt=gmt, rank_metric=dat2, permutation_type=self.permutation_type)

        # reorder datarame for heatmap
        self._heatmat(df=dat.loc[dat2.index], classes=cls_vector,
                      pheno_pos=phenoPos, pheno_neg=phenoNeg)
        # Plotting
        if not self._noplot:
            self._plotting(rank_metric=dat2, results=self.results,
                           graph_num=self.graph_num, outdir=self.outdir,
                           figsize=self.figsize, format=self.format,
                           pheno_pos=phenoPos, pheno_neg=phenoNeg)

        self._logger.info("Congratulations. GSEApy ran successfully.................\n")
        if self._outdir is None:
            self._tmpdir.cleanup()

        return

class Prerank(GSEAbase):
    """GSEA prerank tool"""
    def __init__(self, rnk, gene_sets, outdir='GSEA_prerank',
                 pheno_pos='Pos', pheno_neg='Neg', min_size=15, max_size=500,
                 permutation_num=1000, weighted_score_type=1,
                 ascending=False, processes=1, figsize=(6.5,6), format='pdf',
                 graph_num=20, no_plot=False, seed=None, verbose=False):

        self.rnk =rnk
        self.gene_sets=gene_sets
        self.outdir=outdir
        self.pheno_pos=pheno_pos
        self.pheno_neg=pheno_neg
        self.min_size=min_size
        self.max_size=max_size
        self.permutation_num=int(permutation_num) if int(permutation_num) > 0 else 0
        self.weighted_score_type=weighted_score_type
        self.ascending=ascending
        self.figsize=figsize
        self.format=format
        self.graph_num=int(graph_num)
        self.seed=seed
        self.verbose=bool(verbose)
        self.ranking=None
        self.module='prerank'
        self._processes=processes
        self._noplot=no_plot
        # init logger
        logfile = self.prepare_outdir()
        self._logger = log_init(outlog=logfile,
                                log_level=logging.INFO if self.verbose else logging.WARNING)

    def run(self):
        """GSEA prerank workflow"""

        assert self.min_size <= self.max_size

        # parsing rankings
        dat2 = self._load_ranking(self.rnk)
        assert len(dat2) > 1

        # cpu numbers
        self._set_cores()
        # Start Analysis
        self._logger.info("Parsing data files for GSEA.............................")
        # filtering out gene sets and build gene sets dictionary
        gmt = self.load_gmt(gene_list=dat2.index.values, gmt=self.gene_sets)

        self._logger.info("%04d gene_sets used for further statistical testing....."% len(gmt))
        self._logger.info("Start to run GSEA...Might take a while..................")
        # compute ES, NES, pval, FDR, RES
        gsea_results, hit_ind,rank_ES, subsets = gsea_compute(data=dat2, n=self.permutation_num, gmt=gmt,
                                                              weighted_score_type=self.weighted_score_type,
                                                              permutation_type='gene_set', method=None,
                                                              pheno_pos=self.pheno_pos, pheno_neg=self.pheno_neg,
                                                              classes=None, ascending=self.ascending,
                                                              processes=self._processes, seed=self.seed)
        self._logger.info("Start to generate gseapy reports, and produce figures...")
        res_zip = zip(subsets, list(gsea_results), hit_ind, rank_ES)
        self._save_results(zipdata=res_zip, outdir=self.outdir, module=self.module,
                                   gmt=gmt, rank_metric=dat2, permutation_type="gene_sets")

        # Plotting
        if not self._noplot:
            self._plotting(rank_metric=dat2, results=self.results,
                           graph_num=self.graph_num, outdir=self.outdir,
                           figsize=self.figsize, format=self.format,
                           pheno_pos=self.pheno_pos, pheno_neg=self.pheno_neg)

        self._logger.info("Congratulations. GSEApy runs successfully................\n")
        if self._outdir is None:
            self._tmpdir.cleanup()

        return

class SingleSampleGSEA(GSEAbase):
    """GSEA extension: single sample GSEA"""
    def __init__(self, data, gene_sets, outdir="GSEA_SingleSample", sample_norm_method='rank',
                 min_size=15, max_size=2000, permutation_num=0, weighted_score_type=0.25,
                 scale=True, ascending=False, processes=1, figsize=(7,6), format='pdf',
                 graph_num=20, no_plot=False, seed=None, verbose=False):
        self.data=data
        self.gene_sets=gene_sets
        self.outdir=outdir
        self.sample_norm_method=sample_norm_method
        self.weighted_score_type=weighted_score_type
        self.scale = scale
        self.min_size=min_size
        self.max_size=max_size
        self.permutation_num=int(permutation_num) if int(permutation_num) > 0 else 0
        self.ascending=ascending
        self.figsize=figsize
        self.format=format
        self.graph_num=int(graph_num)
        self.seed=seed
        self.verbose=bool(verbose)
        self.ranking=None
        self.module='ssgsea'
        self._processes=processes
        self._noplot=no_plot
        # init logger

        logfile = self.prepare_outdir()
        self._logger = log_init(outlog=logfile,
                                log_level=logging.INFO if self.verbose else logging.WARNING)

    def corplot(self):
        """NES Correlation plot
        TODO
        """
    def setplot(self):
        """ranked genes' location plot
        TODO
        """

    def load_data(self):
        #load data
        exprs = self.data
        if isinstance(exprs, pd.DataFrame):
            rank_metric = exprs.copy()
            # handle dataframe with gene_name as index.
            self._logger.debug("Input data is a DataFrame with gene names")
            # handle index is not gene_names
            if rank_metric.index.dtype != 'O':
                rank_metric.set_index(keys=rank_metric.columns[0], inplace=True)

            rank_metric = rank_metric.select_dtypes(include=[np.number])
        elif isinstance(exprs, pd.Series):
            # change to DataFrame
            self._logger.debug("Input data is a Series with gene names")
            rank_metric = pd.DataFrame(exprs)
            # rename col if name attr is none
            rank_metric.columns = ["sample1"]
        elif os.path.isfile(exprs):
            # GCT input format?
            if exprs.endswith("gct"):
                rank_metric = pd.read_csv(exprs, skiprows=1, comment='#', index_col=0,sep="\t")
            else:
                # just txt file like input
                rank_metric = pd.read_csv(exprs, comment='#', index_col=0,sep="\t")
                if rank_metric.shape[1] ==1:
                    # rnk file like input
                    rank_metric = pd.read_csv(exprs, header=None, comment='#',
                                                names=['sample1'], index_col=0,sep="\t")
            # select numbers
            rank_metric = rank_metric.select_dtypes(include=[np.number])
        else:
            raise Exception('Error parsing gene ranking values!')

        if rank_metric.index.duplicated().sum() > 0:
            self._logger.warning("Warning: dropping duplicated gene names, only keep the first values")
            rank_metric = rank_metric.loc[rank_metric.index.drop_duplicates(keep='first')]
            rank_metric = rank_metric.loc[rank_metric.index.dropna()]
        if rank_metric.isnull().any().sum() > 0:
            self._logger.warning("Warning: Input data contains NA, filled NA with 0")
            rank_metric = rank_metric.fillna(0)

        return rank_metric

    def norm_samples(self, dat):
        """normalization samples
           see here: http://rowley.mit.edu/caw_web/ssGSEAProjection/ssGSEAProjection.Library.R
        """

        if self.sample_norm_method == 'rank':
            data = dat.rank(axis=0, method='average', na_option='bottom')
            data = 10000*data / data.shape[0]
        elif self.sample_norm_method == 'log_rank':
            data = dat.rank(axis=0, method='average', na_option='bottom')
            data = log(10000*data / data.shape[0] + exp(1))
        elif self.sample_norm_method == 'log':
            dat[dat < 1] = 1
            data = log(dat + exp(1))
        elif self.sample_norm_method == 'custom':
            self._logger.info("Use custom rank metric for ssGSEA")
            data = dat
        else:
            sys.stderr.write("No supported method: %s"%self.sample_norm_method)
            sys.exit(0)

        return data

    def run(self):
        """run entry"""
        self._logger.info("Parsing data files for ssGSEA...........................")
        # load data
        data = self.load_data()
        # normalized samples, and rank
        normdat = self.norm_samples(data)
        # filtering out gene sets and build gene sets dictionary
        gmt = self.load_gmt(gene_list=normdat.index.values, gmt=self.gene_sets)
        self._logger.info("%04d gene_sets used for further statistical testing....."% len(gmt))
        # set cpu numbers
        self._set_cores()
        # start analsis
        self._logger.info("Start to run ssGSEA...Might take a while................")
        if self.permutation_num == 0 :
            # ssGSEA without permutation
            self.runSamples(df=normdat, gmt=gmt)
        else:
            # run permutation procedure and calculate pvals, fdrs
            self._logger.warning("run ssGSEA with permutation procedure, don't use these part of results for publication.")
            self.runSamplesPermu(df=normdat, gmt=gmt)
        # clean up all outputs if _outdir is None
        if self._outdir is None:
            self._tmpdir.cleanup()

    def runSamplesPermu(self, df, gmt=None):
        """Single Sample GSEA workflow with permutation procedure"""

        assert self.min_size <= self.max_size
        mkdirs(self.outdir)
        self.resultsOnSamples = OrderedDict()
        outdir = self.outdir
        # iter throught each sample
        for name, ser in df.iteritems():
            self.outdir = os.path.join(outdir, str(name))
            self._logger.info("Run Sample: %s " % name)
            mkdirs(self.outdir)
            # sort ranking values from high to low or reverse
            dat2 = ser.sort_values(ascending=self.ascending)
            # reset integer index, or caused unwanted problems
            # df.reset_index(drop=True, inplace=True)

            # compute ES, NES, pval, FDR, RES
            gsea_results, hit_ind,rank_ES, subsets = gsea_compute(data=dat2, n=self.permutation_num, gmt=gmt,
                                                                  weighted_score_type=self.weighted_score_type,
                                                                  permutation_type='gene_set', method=None,
                                                                  pheno_pos='', pheno_neg='',
                                                                  classes=None, ascending=self.ascending,
                                                                  processes=self._processes,
                                                                  seed=self.seed, single=True, scale=self.scale)

            # write file
            res_zip = zip(subsets, list(gsea_results), hit_ind, rank_ES)
            self._save_results(zipdata=res_zip, outdir=self.outdir, module=self.module,
                                       gmt=gmt, rank_metric=dat2, permutation_type="gene_sets")
            self.resultsOnSamples[name] = self.res2d.es
            # plotting
            if self._noplot: continue
            self._logger.info("Plotting Sample: %s \n" % name)
            self._plotting(rank_metric=dat2, results=self.results,
                           graph_num=self.graph_num, outdir=self.outdir,
                           figsize=self.figsize, format=self.format)

        # save es, nes to file
        self._save(outdir)

        return

    def runSamples(self, df, gmt=None):
        """Single Sample GSEA workflow.
           multiprocessing utility on samples.
        """

        # df.index.values are gene_names
        # Save each sample results to odict
        self.resultsOnSamples = OrderedDict()
        outdir = self.outdir
        # run ssgsea for gct expression matrix
        #multi-threading
        subsets = sorted(gmt.keys())
        tempes=[]
        names=[]
        rankings=[]
        pool = Pool(processes=self._processes)
        for name, ser in df.iteritems():
            #prepare input
            dat = ser.sort_values(ascending=self.ascending)
            rankings.append(dat)
            names.append(name)
            genes_sorted, cor_vec = dat.index.values, dat.values
            rs = np.random.RandomState(self.seed)
            # apply_async
            tempes.append(pool.apply_async(enrichment_score_tensor,
                                           args=(genes_sorted, cor_vec, gmt,
                                               self.weighted_score_type,
                                               self.permutation_num, rs, True,
                                               self.scale)))
        pool.close()
        pool.join()
        # save results and plotting
        for i, temp in enumerate(tempes):
            name, rnk = names[i], rankings[i]
            self._logger.info("Calculate Enrichment Score for Sample: %s "%name)
            es, esnull, hit_ind, RES = temp.get()
            # create results subdir
            self.outdir= os.path.join(outdir, str(name))
            mkdirs(self.outdir)
            # save results
            self.resultsOnSamples[name] = pd.Series(data=es, index=subsets, name=name)
            # plotting
            if self._noplot: continue
            self._logger.info("Plotting Sample: %s \n" % name)
            for i, term in enumerate(subsets):
                term = term.replace('/','_').replace(":","_")
                outfile = '{0}/{1}.{2}.{3}'.format(self.outdir, term, self.module, self.format)
                gseaplot(rank_metric=rnk, term=term,
                         hits_indices=hit_ind[i], nes=es[i], pval=1, fdr=1,
                         RES=RES[i], pheno_pos='', pheno_neg='',
                         figsize=self.figsize, ofname=outfile)
        # save es, nes to file
        self._save(outdir)

        return

    def _save(self, outdir):
        """save es and stats"""
        # save raw ES to one csv file
        samplesRawES = pd.DataFrame(self.resultsOnSamples)
        samplesRawES.index.name = 'Term|ES'
        # normalize enrichment scores by using the entire data set, as indicated
        # by Barbie et al., 2009, online methods, pg. 2
        samplesNES = samplesRawES / (samplesRawES.values.max() - samplesRawES.values.min())
        samplesNES = samplesNES.copy()
        samplesNES.index.rename('Term|NES', inplace=True)
        self.res2d = samplesNES
        self._logger.info("Congratulations. GSEApy runs successfully................\n")
        if self._outdir is None: return
        # write es
        outESfile = os.path.join(outdir, "gseapy.samples.raw.es.txt")
        with open(outESfile, 'a') as f:
            if self.scale:
                f.write('# scale the enrichment scores by number of genes in the gene sets\n')
                f.write('# this normalization has not effects on the final NES, ' + \
                        'as indicated by Barbie et al., 2009, online methods, pg. 2\n')
            else:
                f.write('# raw enrichment scores of all data\n')
                f.write('# no scale es by numbers of genes in the gene sets\n')
            samplesRawES.to_csv(f, sep='\t')

        outNESfile = os.path.join(outdir, "gseapy.samples.normalized.es.txt")
        with open(outNESfile, 'a') as f:
            f.write('# normalize enrichment scores by using the entire data set\n')
            f.write('# as indicated by Barbie et al., 2009, online methods, pg. 2\n')
            samplesNES.to_csv(f, sep='\t')
        return

class Replot(GSEAbase):
    """To reproduce GSEA desktop output results."""
    def __init__(self, indir, outdir='GSEApy_Replot', weighted_score_type=1,
                  min_size=3, max_size=1000, figsize=(6.5,6), graph_num=20, format='pdf', verbose=False):
        self.indir=indir
        self.outdir=outdir
        self.weighted_score_type=weighted_score_type
        self.min_size=min_size
        self.max_size=max_size
        self.figsize=figsize
        self.fignum=int(graph_num)
        self.format=format
        self.verbose=bool(verbose)
        self.module='replot'
        self.gene_sets=None
        self.ascending=False
        # init logger
        mkdirs(self.outdir)
        outlog = os.path.join(self.outdir,"gseapy.%s.%s.log"%(self.module,"run"))
        self._logger = log_init(outlog=outlog,
                                log_level=logging.INFO if self.verbose else logging.WARNING)
    def run(self):
        """main replot function"""
        assert self.min_size <= self.max_size
        assert self.fignum > 0
        import glob
        from bs4 import BeautifulSoup

        # parsing files.......
        try:
            results_path = glob.glob(self.indir+'*/edb/results.edb')[0]
            rank_path =  glob.glob(self.indir+'*/edb/*.rnk')[0]
            gene_set_path =  glob.glob(self.indir+'*/edb/gene_sets.gmt')[0]
        except IndexError as e:
            sys.stderr.write("Could not locate GSEA files in the given directory!")
            sys.exit(1)
        # extract sample names from .cls file
        cls_path = glob.glob(self.indir+'*/edb/*.cls')
        if cls_path:
            pos, neg, classes = gsea_cls_parser(cls_path[0])
        else:
            # logic for prerank results
            pos, neg = '',''
        # start reploting
        self.gene_sets=gene_set_path
        # obtain gene sets
        gene_set_dict = self.parse_gmt(gmt=gene_set_path)
        # obtain rank_metrics
        rank_metric = self._load_ranking(rank_path)
        correl_vector =  rank_metric.values
        gene_list = rank_metric.index.values
        # extract each enriment term in the results.edb files and plot.
        database = BeautifulSoup(open(results_path), features='xml')
        length = len(database.findAll('DTG'))
        fig_num = self.fignum if self.fignum <= length else length
        for idx in range(fig_num):
            # extract statistical resutls from results.edb file
            enrich_term, hit_ind, nes, pval, fdr= gsea_edb_parser(results_path, index=idx)
            gene_set = gene_set_dict.get(enrich_term)
            # calculate enrichment score
            RES = enrichment_score(gene_list=gene_list,
                                   correl_vector=correl_vector,
                                   gene_set=gene_set,
                                   weighted_score_type=self.weighted_score_type,
                                   nperm=0)[-1]
            # plotting
            term = enrich_term.replace('/','_').replace(":","_")
            outfile = '{0}/{1}.{2}.{3}'.format(self.outdir, term, self.module, self.format)
            gseaplot(rank_metric=rank_metric, term=enrich_term,
                         hits_indices=hit_ind, nes=nes, pval=pval, fdr=fdr,
                         RES=RES, pheno_pos=pos, pheno_neg=neg,
                         figsize=self.figsize, ofname=outfile)

        self._logger.info("Congratulations! Your plots have been reproduced successfully!\n")

def gsea(data, gene_sets, cls, outdir='GSEA_', min_size=15, max_size=500, permutation_num=1000,
          weighted_score_type=1,permutation_type='gene_set', method='log2_ratio_of_classes',
	      ascending=False, processes=1, figsize=(6.5,6), format='pdf',
          graph_num=20, no_plot=False, seed=None, verbose=False):
    """ Run Gene Set Enrichment Analysis.

    :param data: Gene expression data table, Pandas DataFrame, gct file.
    :param gene_sets: Enrichr Library name or .gmt gene sets file or dict of gene sets. Same input with GSEA.
    :param cls: A list or a .cls file format required for GSEA.
    :param str outdir: Results output directory.
    :param int permutation_num: Number of permutations for significance computation. Default: 1000.
    :param str permutation_type: Permutation type, "phenotype" for phenotypes, "gene_set" for genes.
    :param int min_size: Minimum allowed number of genes from gene set also the data set. Default: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Default: 500.
    :param float weighted_score_type: Refer to :func:`algorithm.enrichment_score`. Default:1.
    :param method:  The method used to calculate a correlation or ranking. Default: 'log2_ratio_of_classes'.
                   Others methods are:

                   1. 'signal_to_noise'

                      You must have at least three samples for each phenotype to use this metric.
                      The larger the signal-to-noise ratio, the larger the differences of the means (scaled by the standard deviations);
                      that is, the more distinct the gene expression is in each phenotype and the more the gene acts as a “class marker.”

                   2. 't_test'

                      Uses the difference of means scaled by the standard deviation and number of samples.
                      Note: You must have at least three samples for each phenotype to use this metric.
                      The larger the tTest ratio, the more distinct the gene expression is in each phenotype
                      and the more the gene acts as a “class marker.”

                   3. 'ratio_of_classes' (also referred to as fold change).

                      Uses the ratio of class means to calculate fold change for natural scale data.

                   4. 'diff_of_classes'

                      Uses the difference of class means to calculate fold change for nature scale data

                   5. 'log2_ratio_of_classes'

                      Uses the log2 ratio of class means to calculate fold change for natural scale data.
                      This is the recommended statistic for calculating fold change for log scale data.

    :param bool ascending: Sorting order of rankings. Default: False.
    :param int processes: Number of Processes you are going to use. Default: 1.
    :param list figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [6.5,6].
    :param str format: Matplotlib figure format. Default: 'pdf'.
    :param int graph_num: Plot graphs for top sets of each phenotype.
    :param bool no_plot: If equals to True, no figure will be drawn. Default: False.
    :param seed: Random seed. expect an integer. Default:None.
    :param bool verbose: Bool, increase output verbosity, print out progress of your job, Default: False.

    :return: Return a GSEA obj. All results store to a dictionary, obj.results,
             where contains::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set
                 |  ledge_genes: leading edge genes}

    """
    gs = GSEA(data, gene_sets, cls, outdir, min_size, max_size, permutation_num,
              weighted_score_type, permutation_type, method, ascending, processes,
               figsize, format, graph_num, no_plot, seed, verbose)
    gs.run()

    return gs

def ssgsea(data, gene_sets, outdir="ssGSEA_", sample_norm_method='rank', min_size=15, max_size=2000,
           permutation_num=0, weighted_score_type=0.25, scale=True, ascending=False, processes=1,
           figsize=(7,6), format='pdf', graph_num=20, no_plot=False, seed=None, verbose=False):
    """Run Gene Set Enrichment Analysis with single sample GSEA tool

    :param data: Expression table, pd.Series, pd.DataFrame, GCT file, or .rnk file format.
    :param gene_sets: Enrichr Library name or .gmt gene sets file or dict of gene sets. Same input with GSEA.
    :param outdir: Results output directory.
    :param str sample_norm_method: "Sample normalization method. Choose from {'rank', 'log', 'log_rank'}. Default: rank.

               1. 'rank': Rank your expression data, and transform by 10000*rank_dat/gene_numbers
               2. 'log' : Do not rank, but transform data by log(data + exp(1)), while data = data[data<1] =1.
               3. 'log_rank': Rank your expression data, and transform by log(10000*rank_dat/gene_numbers+ exp(1))
               4. 'custom': Do nothing, and use your own rank value to calculate enrichment score.

    see here: https://github.com/GSEA-MSigDB/ssGSEAProjection-gpmodule/blob/master/src/ssGSEAProjection.Library.R, line 86

    :param int min_size: Minimum allowed number of genes from gene set also the data set. Default: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Default: 2000.
    :param int permutation_num: Number of permutations for significance computation. Default: 0.
    :param str weighted_score_type: Refer to :func:`algorithm.enrichment_score`. Default:0.25.
    :param bool scale: If True, normalize the scores by number of genes in the gene sets.
    :param bool ascending: Sorting order of rankings. Default: False.
    :param int processes: Number of Processes you are going to use. Default: 1.
    :param list figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [7,6].
    :param str format: Matplotlib figure format. Default: 'pdf'.
    :param int graph_num: Plot graphs for top sets of each phenotype.
    :param bool no_plot: If equals to True, no figure will be drawn. Default: False.
    :param seed: Random seed. expect an integer. Default:None.
    :param bool verbose: Bool, increase output verbosity, print out progress of your job, Default: False.

    :return: Return a ssGSEA obj.
             All results store to  a dictionary, access enrichment score by obj.resultsOnSamples,
             and normalized enrichment score by obj.res2d.
             if permutation_num > 0, additional results contain::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set
                 |  ledge_genes: leading edge genes, if permutation_num >0}

    """

    ss = SingleSampleGSEA(data, gene_sets, outdir, sample_norm_method, min_size, max_size,
                          permutation_num, weighted_score_type, scale, ascending,
                          processes, figsize, format, graph_num, no_plot, seed, verbose)
    ss.run()
    return ss

def prerank(rnk, gene_sets, outdir='GSEA_Prerank', pheno_pos='Pos', pheno_neg='Neg',
            min_size=15, max_size=500, permutation_num=1000, weighted_score_type=1,
            ascending=False, processes=1, figsize=(6.5,6), format='pdf',
            graph_num=20, no_plot=False, seed=None, verbose=False):
    """ Run Gene Set Enrichment Analysis with pre-ranked correlation defined by user.

    :param rnk: pre-ranked correlation table or pandas DataFrame. Same input with ``GSEA`` .rnk file.
    :param gene_sets: Enrichr Library name or .gmt gene sets file or dict of gene sets. Same input with GSEA.
    :param outdir: results output directory.
    :param int permutation_num: Number of permutations for significance computation. Default: 1000.
    :param int min_size: Minimum allowed number of genes from gene set also the data set. Default: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Defaults: 500.
    :param str weighted_score_type: Refer to :func:`algorithm.enrichment_score`. Default:1.
    :param bool ascending: Sorting order of rankings. Default: False.
    :param int processes: Number of Processes you are going to use. Default: 1.
    :param list figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [6.5,6].
    :param str format: Matplotlib figure format. Default: 'pdf'.
    :param int graph_num: Plot graphs for top sets of each phenotype.
    :param bool no_plot: If equals to True, no figure will be drawn. Default: False.
    :param seed: Random seed. expect an integer. Default:None.
    :param bool verbose: Bool, increase output verbosity, print out progress of your job, Default: False.

    :return: Return a Prerank obj. All results store to  a dictionary, obj.results,
             where contains::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set
                 |  ledge_genes: leading edge genes}

    """
    pre = Prerank(rnk, gene_sets, outdir, pheno_pos, pheno_neg,
                  min_size, max_size, permutation_num, weighted_score_type,
                  ascending, processes, figsize, format, graph_num, no_plot, seed, verbose)
    pre.run()
    return pre

def replot(indir, outdir='GSEA_Replot', weighted_score_type=1,
           min_size=3, max_size=1000, figsize=(6.5,6), graph_num=20, format='pdf', verbose=False):
    """The main function to reproduce GSEA desktop outputs.

    :param indir: GSEA desktop results directory. In the sub folder, you must contain edb file folder.
    :param outdir: Output directory.
    :param float weighted_score_type: weighted score type. choose from {0,1,1.5,2}. Default: 1.
    :param list figsize: Matplotlib output figure figsize. Default: [6.5,6].
    :param str format: Matplotlib output figure format. Default: 'pdf'.
    :param int min_size: Min size of input genes presented in Gene Sets. Default: 3.
    :param int max_size: Max size of input genes presented in Gene Sets. Default: 5000.
                     You are not encouraged to use min_size, or max_size argument in :func:`replot` function.
                     Because gmt file has already been filtered.
    :param verbose: Bool, increase output verbosity, print out progress of your job, Default: False.

    :return: Generate new figures with selected figure format. Default: 'pdf'.

    """
    rep = Replot(indir, outdir, weighted_score_type,
                 min_size, max_size, figsize, graph_num, format, verbose)
    rep.run()

    return


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# see: http://amp.pharm.mssm.edu/Enrichr/help#api for API docs

import sys, json, os, logging
import requests
import pandas as pd
from io import StringIO
from collections import OrderedDict
from pkg_resources import resource_filename
from time import sleep
from tempfile import TemporaryDirectory
from numpy import isscalar
from gseapy.plot import barplot
from gseapy.parser import Biomart
from gseapy.utils import *
from gseapy.stats import calc_pvalues, multiple_testing_correction

class Enrichr(object):
    """Enrichr API"""
    def __init__(self, gene_list, gene_sets, organism='human', descriptions='',
                 outdir='Enrichr', cutoff=0.05, background='hsapiens_gene_ensembl',
                 format='pdf', figsize=(6.5,6), top_term=10, no_plot=False,
                 verbose=False):

        self.gene_list = gene_list
        self.gene_sets = gene_sets
        self.descriptions = str(descriptions)
        self.outdir = outdir
        self.cutoff = cutoff
        self.format = format
        self.figsize = figsize
        self.__top_term = int(top_term)
        self.__no_plot = no_plot
        self.verbose = bool(verbose)
        self.module = "enrichr"
        self.res2d = None
        self._processes = 1
        self.background = background
        self._bg = None
        self.organism = organism
        self._organism = None
        # init logger
        logfile = self.prepare_outdir()
        self._logger = log_init(outlog=logfile,
                                log_level=logging.INFO if self.verbose else logging.WARNING)

    def prepare_outdir(self):
        """create temp directory."""
        self._outdir = self.outdir
        if self._outdir is None:
            self._tmpdir = TemporaryDirectory()
            self.outdir = self._tmpdir.name
        elif isinstance(self.outdir, str):
            mkdirs(self.outdir)
        else:
            raise Exception("Error parsing outdir: %s"%type(self.outdir))

        # handle gene_sets
        logfile = os.path.join(self.outdir, "gseapy.%s.%s.log" % (self.module, self.descriptions))
        return logfile

    def parse_genesets(self):
        """parse gene_sets input file type"""

        enrichr_library = self.get_libraries()
        if isinstance(self.gene_sets, list):
            gss = self.gene_sets
        elif isinstance(self.gene_sets, str):
            gss = [ g.strip() for g in self.gene_sets.strip().split(",") ]
        elif isinstance(self.gene_sets, dict):
            gss = [self.gene_sets]
        else:
            raise Exception("Error parsing enrichr libraries, please provided corrected one")

        # gss: a list contain .gmt, dict, enrichr_liraries.
        # now, convert .gmt to dict
        gss_exist = []
        for g in gss:
            if isinstance(g, dict):
                gss_exist.append(g)
                continue

            if isinstance(g, str):
                if g in enrichr_library:
                    gss_exist.append(g)
                    continue
                if g.lower().endswith(".gmt") and os.path.exists(g):
                    self._logger.info("User Defined gene sets is given: %s"%g)
                    with open(g) as genesets:
                        g_dict = { line.strip().split("\t")[0]: line.strip().split("\t")[2:]
                                        for line in genesets.readlines() }
                    gss_exist.append(g_dict)
        return gss_exist

    def parse_genelists(self):
        """parse gene list"""
        if isinstance(self.gene_list, list):
            genes = self.gene_list
        elif isinstance(self.gene_list, pd.DataFrame):
            # input type is bed file
            if self.gene_list.shape[1] >=3:
                genes= self.gene_list.iloc[:,:3].apply(lambda x: "\t".join([str(i) for i in x]), axis=1).tolist()
            # input type with weight values
            elif self.gene_list.shape[1] == 2:
               genes= self.gene_list.apply(lambda x: ",".join([str(i) for i in x]), axis=1).tolist()
            else:
               genes = self.gene_list.squeeze().tolist()
        elif isinstance(self.gene_list, pd.Series):
            genes = self.gene_list.squeeze().tolist()
        else:
            # get gene lists or bed file, or gene list with weighted values.
            genes=[]
            with open(self.gene_list) as f:
                for gene in f:
                    genes.append(gene.strip())

        self._isezid = all(map(self._is_entrez_id, genes))
        if self._isezid:
            self._gls = set(map(int, self._gls))
        else:
            self._gls = genes

        return '\n'.join(genes)

    def send_genes(self, gene_list, url):
        """ send gene list to enrichr server"""
        payload = {
          'list': (None, gene_list),
          'description': (None, self.descriptions)
           }
        # response
        response = requests.post(url, files=payload)
        if not response.ok:
            raise Exception('Error analyzing gene list')
        sleep(1)
        job_id = json.loads(response.text)

        return job_id

    def check_genes(self, gene_list, usr_list_id):
        '''
        Compare the genes sent and received to get successfully recognized genes
        '''
        response = requests.get('http://amp.pharm.mssm.edu/Enrichr/view?userListId=%s' % usr_list_id)
        if not response.ok:
            raise Exception('Error getting gene list back')
        returnedL = json.loads(response.text)["genes"]
        returnedN = sum([1 for gene in gene_list if gene in returnedL])
        self._logger.info('{} genes successfully recognized by Enrichr'.format(returnedN))

    def get_results(self, gene_list):
        """Enrichr API"""
        ADDLIST_URL = 'http://amp.pharm.mssm.edu/%sEnrichr/addList'%self._organism
        job_id = self.send_genes(gene_list, ADDLIST_URL)
        user_list_id = job_id['userListId']

        RESULTS_URL = 'http://amp.pharm.mssm.edu/%sEnrichr/export'%self._organism
        query_string = '?userListId=%s&filename=%s&backgroundType=%s'
        # set max retries num =5
        s = retry(num=5)
        filename = "%s.%s.reports" % (self._gs, self.descriptions)
        url = RESULTS_URL + query_string % (user_list_id, filename, self._gs)
        response = s.get(url, stream=True, timeout=None)
        # response = requests.get(RESULTS_URL + query_string % (user_list_id, gene_set))
        sleep(1)
        res = pd.read_csv(StringIO(response.content.decode('utf-8')),sep="\t")
        return [job_id['shortId'], res]

    def _is_entrez_id(self, idx):
        try:
            int(idx)
            return True
        except:
            return False

    def get_libraries(self,):
        """return active enrichr library name. Official API """
        lib_url='http://amp.pharm.mssm.edu/%sEnrichr/datasetStatistics'%self._organism
        libs_json = json.loads(requests.get(lib_url).text)
        libs = [lib['libraryName'] for lib in libs_json['statistics']]
        return sorted(libs)

    def get_background(self):
        """get background gene"""

        # input is a file
        if os.path.isfile(self.background):
            with open(self.background) as b:
                bg2 = b.readlines()
            bg = [g.strip() for g in bg2]
            return set(bg)

        # package included data
        DB_FILE = resource_filename("gseapy", "data/{}.background.genes.txt".format(self.background))
        filename = os.path.join(DEFAULT_CACHE_PATH, "{}.background.genes.txt".format(self.background))
        if os.path.exists(filename):
            df = pd.read_csv(filename,sep="\t")
        elif os.path.exists(DB_FILE):
            df = pd.read_csv(DB_FILE,sep="\t")
        else:
            # background is a biomart database name
            self._logger.warning("Downloading %s for the first time. It might take a couple of miniutes."%self.background)
            bm = Biomart()
            df = bm.query(dataset=self.background)
            df.dropna(subset=['go_id'], inplace=True)
        self._logger.info("using all annotated genes with GO_ID as background genes")
        df.dropna(subset=['entrezgene'], inplace=True)
        # input id type: entrez or gene_name
        if self._isezid:
            bg = df['entrezgene'].astype(int)
        else:
            bg = df['external_gene_name']

        return set(bg)

    def get_organism(self):
        """Select Enrichr organism from below:

           Human & Mouse: H. sapiens & M. musculus
           Fly: D. melanogaster
           Yeast: S. cerevisiae
           Worm: C. elegans
           Fish: D. rerio

        """

        organism = {'default': ['', 'hs', 'mm', 'human','mouse',
                                'homo sapiens', 'mus musculus',
                                'h. sapiens', 'm. musculus'],
                    'Fly': ['fly', 'd. melanogaster', 'drosophila melanogaster'],
                    'Yeast': ['yeast', 's. cerevisiae', 'saccharomyces cerevisiae'],
                    'Worm': ['worm', 'c. elegans', 'caenorhabditis elegans', 'nematode'],
                    'Fish': ['fish', 'd. rerio', 'danio rerio', 'zebrafish']
                 }

        for k, v in organism.items():
            if self.organism.lower() in v :
                self._organism = k

        if self._organism is None:
            raise Exception("No supported organism found !!!")

        if self._organism == 'default':
            self._organism = ''
        return

    def enrich(self, gmt):
        """use local mode

        p = p-value computed using the Fisher exact test (Hypergeometric test)

        Not implemented here:

            combine score = log(p)·z

        see here: http://amp.pharm.mssm.edu/Enrichr/help#background&q=4

        columns contain:

            Term Overlap P-value Adjusted_P-value Genes

        """
        if isscalar(self.background):
            if isinstance(self.background, int) or self.background.isdigit():
                self._bg = int(self.background)
            elif isinstance(self.background, str):
                # self.background = set(reduce(lambda x,y: x+y, gmt.values(),[]))
                self._bg = self.get_background()
                self._logger.info("Background: found %s genes"%(len(self._bg)))
            else:
                raise Exception("Unsupported background data type")
        else:
            # handle array object: nd.array, list, tuple, set, Series
            try:
                it = iter(self.background)
                self._bg = set(self.background)
            except TypeError:
                self._logger.error("Unsupported background data type")
        # statistical testing
        hgtest = list(calc_pvalues(query=self._gls, gene_sets=gmt,
                                   background=self._bg))
        if len(hgtest) > 0:
            terms, pvals, olsz, gsetsz, genes = hgtest
            fdrs, rej = multiple_testing_correction(ps = pvals,
                                                    alpha=self.cutoff,
                                                    method='benjamini-hochberg')
            # save to a dataframe
            odict = OrderedDict()
            odict['Term'] = terms
            odict['Overlap'] = list(map(lambda h,g: "%s/%s"%(h, g), olsz, gsetsz))
            odict['P-value'] = pvals
            odict['Adjusted P-value'] = fdrs
            # odict['Reject (FDR< %s)'%self.cutoff ] = rej
            odict['Genes'] = [";".join(g) for g in genes]
            res = pd.DataFrame(odict)
            return res
        return

    def run(self):
        """run enrichr for one sample gene list but multi-libraries"""

        # set organism
        self.get_organism()
        # read input file
        genes_list = self.parse_genelists()
        gss = self.parse_genesets()
        # if gmt
        self._logger.info("Connecting to Enrichr Server to get latest library names")
        if len(gss) < 1:
            sys.stderr.write("Not validated Enrichr library name provided\n")
            sys.stdout.write("Hint: use get_library_name() to view full list of supported names")
            sys.exit(1)
        self.results = pd.DataFrame()

        for g in gss:
            if isinstance(g, dict):
                ## local mode
                res = self.enrich(g)
                shortID, self._gs = str(id(g)), "CUSTOM%s"%id(g)
                if res is None:
                    self._logger.info("No hits return, for gene set: Custom%s"%shortID)
                    continue
            else:
                ## online mode
                self._gs = str(g)
                self._logger.debug("Start Enrichr using library: %s" % (self._gs))
                self._logger.info('Analysis name: %s, Enrichr Library: %s' % (self.descriptions, self._gs))
                shortID, res = self.get_results(genes_list)
                # Remember gene set library used
            res.insert(0, "Gene_set", self._gs)
            # Append to master dataframe
            self.results = self.results.append(res, ignore_index=True, sort=True)
            self.res2d = res
            if self._outdir is None: continue
            self._logger.info('Save file of enrichment results: Job Id:' + str(shortID))
            outfile = "%s/%s.%s.%s.reports.txt" % (self.outdir, self._gs, self.descriptions, self.module)
            self.res2d.to_csv(outfile, index=False, encoding='utf-8', sep="\t")
            # plotting
            if not self.__no_plot:
                msg = barplot(df=res, cutoff=self.cutoff, figsize=self.figsize,
                              top_term=self.__top_term, color='salmon',
                              title=self._gs,
                              ofname=outfile.replace("txt", self.format))
                if msg is not None : self._logger.warning(msg)
            self._logger.info('Done.\n')
        # clean up tmpdir
        if self._outdir is None: self._tmpdir.cleanup()

        return

def enrichr(gene_list, gene_sets, organism='human', description='',
            outdir='Enrichr', background='hsapiens_gene_ensembl', cutoff=0.05,
            format='pdf', figsize=(8,6), top_term=10, no_plot=False, verbose=False):
    """Enrichr API.

    :param gene_list: Flat file with list of genes, one gene id per row, or a python list object
    :param gene_sets: Enrichr Library to query. Required enrichr library name(s). Separate each name by comma.
    :param organism: Enrichr supported organism. Select from (human, mouse, yeast, fly, fish, worm).
                     see here for details: https://amp.pharm.mssm.edu/modEnrichr
    :param description: name of analysis. optional.
    :param outdir: Output file directory
    :param float cutoff: Adjusted P-value (benjamini-hochberg correction) cutoff. Default: 0.05
    :param int background: BioMart dataset name for retrieving background gene information.
                           This argument only works when gene_sets input is a gmt file or python dict.
                           You could also specify a number by yourself, e.g. total expressed genes number.
                           In this case, you will skip retrieving background infos from biomart.

    Use the code below to see valid background dataset names from BioMart.
    Here are example code:
    >>> from gseapy.parser import Biomart
    >>> bm = Biomart(verbose=False, host="asia.ensembl.org")
    >>> ## view validated marts
    >>> marts = bm.get_marts()
    >>> ## view validated dataset
    >>> datasets = bm.get_datasets(mart='ENSEMBL_MART_ENSEMBL')

    :param str format: Output figure format supported by matplotlib,('pdf','png','eps'...). Default: 'pdf'.
    :param list figsize: Matplotlib figsize, accept a tuple or list, e.g. (width,height). Default: (6.5,6).
    :param bool no_plot: If equals to True, no figure will be drawn. Default: False.
    :param bool verbose: Increase output verbosity, print out progress of your job, Default: False.

    :return: An Enrichr object, which obj.res2d stores your last query, obj.results stores your all queries.

    """
    enr = Enrichr(gene_list, gene_sets, organism, description, outdir,
                  cutoff, background, format, figsize, top_term, no_plot, verbose)
    enr.run()

    return enr


#! python3
""" MeshLabXML creation functions """

import math

from . import FilterScript
from . import util
from . import transform
from . import vert_color
from . import clean
from . import layers

def cube(script, size=1.0, center=False, color=None):
    """Create a cube primitive

    Note that this is made of 6 quads, not triangles
    """

    """# Convert size to list if it isn't already
    if not isinstance(size, list):
        size = list(size)
    # If a single value was supplied use it for all 3 axes
    if len(size) == 1:
        size = [size[0], size[0], size[0]]"""
    size = util.make_list(size, 3)
    if script.ml_version == '1.3.4BETA':
        filter_name = 'Box'
    else:
        filter_name = 'Box/Cube'
    filter_xml = ''.join([
        '  <filter name="{}">\n'.format(filter_name),
        '    <Param name="size" ',
        'value="1.0" ',
        'description="Scale factor" ',
        'type="RichFloat" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Cube', change_layer=True)
    transform.scale(script, value=size)
    # Box is centered on origin at creation
    if not center:
        transform.translate(script, value=[size[0]/2, size[1]/2, size[2]/2])
    if color is not None:
        vert_color.function(script, color=color)
    return None

# Usage: height=(1) (radius=(1)|(radius1=(1) radius2=(1)))|(diameter=(2)|(diameter1=(2) diameter2=(2))) center=(false)
# Note: need to know m_up in order to orient cylinder correctly!
# OpenSCAD cylinder:
#    height This is the height of the cylinder. Default value is 1.
#    radius The radius of both top and bottom ends of the cylinder. Use this parameter if you want plain cylinder. Default value is 1.
#    radius1 This is the radius of the cone on bottom end. Default value is 1.
#    radius2 This is the radius of the cone on top end. Default value is 1.
#    diameter The diameter of both top and bottom ends of the cylinder. Use this parameter if you want plain cylinder. Default value is 1.
#    diameter1 This is the diameter of the cone on bottom end. Default value is 1.
#    diameter2 This is the diameter of the cone on top end. Default value is 1.
# center If true will center the height of the cone/cylinder around the
# origin. Default is false, placing the base of the cylinder or radius1 radius
# of cone at the origin.
def cylinder(script, up='z', height=1.0, radius=None, radius1=None,
             radius2=None, diameter=None, diameter1=None, diameter2=None,
             center=False, cir_segments=32, color=None):
    """Create a cylinder or cone primitive. Usage is based on OpenSCAD.
    # height = height of the cylinder
    # radius1 = radius of the cone on bottom end
    # radius2 = radius of the cone on top end
    # center = If true will center the height of the cone/cylinder around
    # the origin. Default is false, placing the base of the cylinder or radius1
    # radius of cone at the origin.
    #
    cir_segments Number of sides of the polygonal approximation of the cone

    color = specify a color name to apply vertex colors to the newly
    # created mesh
    """
    if radius is not None and diameter is None:
        if radius1 is None and diameter1 is None:
            radius1 = radius
        if radius2 is None and diameter2 is None:
            radius2 = radius
    if diameter is not None:
        if radius1 is None and diameter1 is None:
            radius1 = diameter / 2
        if radius2 is None and diameter2 is None:
            radius2 = diameter / 2
    if diameter1 is not None:
        radius1 = diameter1 / 2
    if diameter2 is not None:
        radius2 = diameter2 / 2
    if radius1 is None:
        radius1 = 1.0
    if radius2 is None:
        radius2 = radius1

    # Cylinder is created centered with Y up
    filter_xml = ''.join([
        '  <filter name="Cone">\n',
        '    <Param name="h" ',
        'value="%s" ' % height,
        'description="Height" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="r0" ',
        'value="%s" ' % radius1,
        'description="Radius 1" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="r1" ',
        'value="%s" ' % radius2,
        'description="Radius 2" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="subdiv" ',
        'value="%d" ' % cir_segments,
        'description="Side" ',
        'type="RichInt" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Cone', change_layer=True)
    if not center:
        transform.translate(script, [0, height / 2, 0])
    if up.lower() == 'z':
        transform.rotate(script, axis='x', angle=90)  # rotate to Z up
    if color is not None:
        vert_color.function(script, color=color)
    return None

def icosphere(script, radius=1.0, diameter=None, subdivisions=3, color=None):
    """create an icosphere mesh

    radius Radius of the sphere
    # subdivisions = Subdivision level; Number of the recursive subdivision of the
    # surface. Default is 3 (a sphere approximation composed by 1280 faces).
    # Admitted values are in the range 0 (an icosahedron) to 8 (a 1.3 MegaTris
    # approximation of a sphere). Formula for number of faces: F=20*4^subdiv
    # color = specify a color name to apply vertex colors to the newly
    # created mesh"""
    if diameter is not None:
        radius = diameter / 2
    filter_xml = ''.join([
        '  <filter name="Sphere">\n',
        '    <Param name="radius" ',
        'value="%s" ' % radius,
        'description="Radius" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="subdiv" ',
        'value="%d" ' % subdivisions,
        'description="Subdiv. Level" ',
        'type="RichInt" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Sphere', change_layer=True)
    if color is not None:
        vert_color.function(script, color=color)
    return None

def sphere_cap(script, angle=1.0, subdivisions=3, color=None):
    """# angle = Angle of the cone subtending the cap. It must be <180 less than 180
    # subdivisions = Subdivision level; Number of the recursive subdivision of the
    # surface. Default is 3 (a sphere approximation composed by 1280 faces).
    # Admitted values are in the range 0 (an icosahedron) to 8 (a 1.3 MegaTris
    # approximation of a sphere). Formula for number of faces: F=20*4^subdivisions
    # color = specify a color name to apply vertex colors to the newly
    # created mesh"""
    filter_xml = ''.join([
        '  <filter name="Sphere Cap">\n',
        '    <Param name="angle" ',
        'value="%s" ' % angle,
        'description="Angle" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="subdiv" ',
        'value="%d" ' % subdivisions,
        'description="Subdiv. Level" ',
        'type="RichInt" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Sphere Cap', change_layer=True)
    if color is not None:
        vert_color.function(script, color=color)
    return None

def torus(script, major_radius=3.0, minor_radius=1.0, inner_diameter=None,
          outer_diameter=None, major_segments=48, minor_segments=12,
          color=None):
    """Create a torus mesh

    Args:
        major_radius (float, (optional)): radius from the origin to the
            center of the cross sections
        minor_radius (float, (optional)): radius of the torus cross
            section
        inner_diameter (float, (optional)): inner diameter of torus. If
            both inner_diameter and outer_diameter are provided then
            these will override major_radius and minor_radius.,
        outer_diameter (float, (optional)): outer diameter of torus. If
            both inner_diameter and outer_diameter are provided then
            these will override major_radius and minor_radius.
        major_segments (int (optional)): number of segments for the main
            ring of the torus
        minor_segments (int (optional)): number of segments for the minor
            ring of the torus
        color (str (optional)): color name to apply vertex colors to the
            newly created mesh

    Returns:
        None

    """
    if inner_diameter is not None and outer_diameter is not None:
        major_radius = (inner_diameter + outer_diameter) / 4
        minor_radius = major_radius - inner_diameter / 2
        # Ref: inner_diameter = 2 * (major_radius - minor_radius)
        # Ref: outer_diameter = 2 * (major_radius + minor_radius)
    filter_xml = ''.join([
        '  <filter name="Torus">\n',
        '    <Param name="hRadius" ',
        'value="%s" ' % major_radius,
        'description="Horizontal Radius" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="vRadius" ',
        'value="%s" ' % minor_radius,
        'description="Vertical Radius" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="hSubdiv" ',
        'value="%d" ' % major_segments,
        'description="Horizontal Subdivision" ',
        'type="RichInt" ',
        '/>\n',
        '    <Param name="vSubdiv" ',
        'value="%d" ' % minor_segments,
        'description="Vertical Subdivision" ',
        'type="RichInt" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Torus', change_layer=True)
    if color is not None:
        vert_color.function(script, color=color)
    return None

def grid(script, size=1.0, x_segments=1, y_segments=1, center=False,
         color=None):
    """2D square/plane/grid created on XY plane

    x_segments # Number of segments in the X direction.
    y_segments # Number of segments in the Y direction.
    center="false" # If true square will be centered on origin;
    otherwise it is place in the positive XY quadrant.

    """
    size = util.make_list(size, 2)
    filter_xml = ''.join([
        '  <filter name="Grid Generator">\n',
        '    <Param name="absScaleX" ',
        'value="{}" '.format(size[0]),
        'description="x scale" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="absScaleY" ',
        'value="{}" '.format(size[1]),
        'description="y scale" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="numVertX" ',
        'value="{:d}" '.format(x_segments + 1),
        'description="num vertices on x" ',
        'type="RichInt" ',
        '/>\n',
        '    <Param name="numVertY" ',
        'value="{:d}" '.format(y_segments + 1),
        'description="num vertices on y" ',
        'type="RichInt" ',
        '/>\n',
        '    <Param name="center" ',
        'value="false" ',
        'description="centered on origin" ',
        'type="RichBool" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Grid Generator', change_layer=True)

    """This is to work around a bug in MeshLab whereby the Grid Generator does not
    create zero values for z. Ref bug #458: https://sourceforge.net/p/meshlab/bugs/458/"""
    transform.vert_function(script, z_func='rint(z)')

    """Note that the "center" parameter in the mlx script does not actually
    center the square, not sure what it is doing. Instead this is set to false,
    which places the plane in the -X,+Y quadrant, and it is translated to the
    appropriate position after creation.
    """
    if center:
        transform.translate(script, value=[size[0]/2, -size[1]/2, 0])
    else:
        transform.translate(script, value=[size[0], 0, 0])
    if color is not None:
        vert_color.function(script, color=color)
    return None

def annulus(script, radius=None, radius1=None, radius2=None, diameter=None,
            diameter1=None, diameter2=None, cir_segments=32, color=None):
    """Create a 2D (surface) circle or annulus
    radius1=1 # Outer radius of the circle
    radius2=0 # Inner radius of the circle (if non-zero it creates an annulus)
    color="" # specify a color name to apply vertex colors to the newly created mesh

    OpenSCAD: parameters: diameter overrides radius, radius1 & radius2 override radius
    """
    if radius is not None and diameter is None:
        if radius1 is None and diameter1 is None:
            radius1 = radius
        if radius2 is None and diameter2 is None:
            radius2 = 0
    if diameter is not None:
        if radius1 is None and diameter1 is None:
            radius1 = diameter / 2
        if radius2 is None and diameter2 is None:
            radius2 = 0
    if diameter1 is not None:
        radius1 = diameter1 / 2
    if diameter2 is not None:
        radius2 = diameter2 / 2
    if radius1 is None:
        radius1 = 1
    if radius2 is None:
        radius2 = 0

    # Circle is created centered on the XY plane
    filter_xml = ''.join([
        '  <filter name="Annulus">\n',
        '    <Param name="externalRadius" ',
        'value="%s" ' % radius1,
        'description="External Radius" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="internalRadius" ',
        'value="%s" ' % radius2,
        'description="Internal Radius" ',
        'type="RichFloat" ',
        '/>\n',
        '    <Param name="sides" ',
        'value="%d" ' % cir_segments,
        'description="Sides" ',
        'type="RichInt" ',
        '/>\n',
        '  </filter>\n'])
    util.write_filter(script, filter_xml)
    if isinstance(script, FilterScript):
        script.add_layer('Annulus', change_layer=True)
    if color is not None:
        vert_color.function(script, color=color)
    return None

def cylinder_open_hires(script, height=1.0, radius=1, diameter=None,
                        cir_segments=48, height_segments=1,
                        invert_normals=False, center=False, color=None):
    """ Creates a round open tube, e.g. a cylinder with no top or bottom.

    Useful if you want to wrap it around and join the open ends together, forming a torus.

    invert_normals (bool (optional)): if True normals point outward; in false normals point inward.
    """
    if diameter is not None:
        radius = diameter / 2

    if center:
        z_translate = -height / 2
    else:
        z_translate = 0.0

    grid(script,
         [2 * math.pi * radius, height],
         x_segments=cir_segments,
         y_segments=height_segments)
    transform.rotate(script, 'x', 90)
    transform.translate(script, [math.pi * radius / 2, 0, z_translate])
    if not invert_normals:
        transform.rotate(script, 'z', 180)
    transform.wrap2cylinder(script, radius)
    clean.merge_vert(script, threshold=0.00002)
    if color is not None:
        vert_color.function(script, color=color)
    return None

def cube_open_hires_old(script, size=1.0, x_segments=1, y_segments=1, z_segments=1,
                    center=False, color=None):
    """ Creates a square open tube, e.g. a box with no top or bottom.

    Useful if you want to wrap it around and join the open ends together, forming a torus.
    """
    """# Convert size to list if it isn't already
    if not isinstance(size, list):
        size = list(size)
    # If a single value was supplied use it for all 3 axes
    if len(size) == 1:
        size = [size[0], size[0], size[0]]"""
    size = util.make_list(size, 3)

    # X sides
    grid(script, [size[0], size[2]],
         x_segments=x_segments,
         y_segments=z_segments)
    transform.rotate(script, 'x', 90)
    #transform.translate(script, [0, 0, -size[2]])
    layers.duplicate(script)
    # Rotate to correct normals
    transform.rotate(script, 'z', 180)
    transform.translate(script, [size[0], size[1], 0])

    # Y sides
    grid(script, [size[2], size[1]],
         x_segments=z_segments,
         y_segments=y_segments)
    transform.rotate(script, 'y', -90)
    #transform.rotate(script, 'z', 90)
    #transform.translate(script, [0, 0, -size[2]])
    layers.duplicate(script)
    # Rotate to correct normals
    transform.rotate(script, 'z', 180)
    transform.translate(script, [size[0], size[1], 0])

    layers.join(script)
    clean.merge_vert(script, threshold=0.00002)
    # normals.fix(script)
    if center:
        transform.translate(script, [-size[0] / 2, -size[1] / 2, -size[2] / 2])
    if color is not None:
        vert_color.function(script, color=color)
    return None

def cube_open_hires(script, size=1.0, x_segments=1, y_segments=1, z_segments=1,
                    center=False, color=None):
    """ Creates a square open tube, e.g. a box with no top or bottom.

    Useful if you want to wrap it around and join the open ends together, forming a torus.
    """
    """# Convert size to list if it isn't already
    if not isinstance(size, list):
        size = list(size)
    # If a single value was supplied use it for all 3 axes
    if len(size) == 1:
        size = [size[0], size[0], size[0]]"""
    size = util.make_list(size, 3)

    # Make big grid and bend
    grid(script, [2*(x_segments + y_segments), z_segments],
         x_segments=2*(x_segments + y_segments),
         y_segments=z_segments)
    transform.rotate(script, 'x', 90)
    # Bend 3 times into a rectangular tube
    if script.ml_version == '1.3.4BETA': # muparser version: 1.3.2
        transform.vert_function(script,
            x_func='if(x>{x_size}, {x_size}, x)'.format(x_size=x_segments),
            y_func='if(x>{x_size}, (x-{x_size}), y)'.format(x_size=x_segments),
            z_func='z')
        transform.vert_function(script,
            x_func='if(y>{y_size}, ({y_size}-y+{x_size}), x)'.format(x_size=x_segments, y_size=y_segments),
            y_func='if(y>{y_size}, {y_size}, y)'.format(y_size=y_segments),
            z_func='z')
        transform.vert_function(script,
            x_func='if(x<0, 0, x)',
            y_func='if(x<0, ({y_size}+x), y)'.format(y_size=y_segments),
            z_func='z')
    else: # muparser version: 2.2.5
        transform.vert_function(script,
            x_func='(x>{x_size} ? {x_size} : x)'.format(x_size=x_segments),
            y_func='(x>{x_size} ? (x-{x_size}) : y)'.format(x_size=x_segments),
            z_func='z')
        transform.vert_function(script,
            x_func='(y>{y_size} ? ({y_size}-y+{x_size}) : x)'.format(x_size=x_segments, y_size=y_segments),
            y_func='(y>{y_size} ? {y_size} : y)'.format(y_size=y_segments),
            z_func='z')
        transform.vert_function(script,
            x_func='(x<0 ? 0 : x)',
            y_func='(x<0 ? ({y_size}+x) : y)'.format(y_size=y_segments),
            z_func='z')
    clean.merge_vert(script, threshold=0.00002)
    transform.scale(script, [size[0]/x_segments, size[1]/y_segments, size[2]/z_segments])
    if center:
        transform.translate(script, [-size[0] / 2, -size[1] / 2, -size[2] / 2])
    if color is not None:
        vert_color.function(script, color=color)
    return None

def plane_hires_edges(script, size=1.0, x_segments=1, y_segments=1,
                      center=False, color=None):
    """ Creates a plane with a specified number of vertices
    on it sides, but no vertices on the interior.

    Currently used to create a simpler bottom for cube_hires.

    """
    size = util.make_list(size, 2)

    grid(script, size=[x_segments + y_segments - 1, 1],
         x_segments=(x_segments + y_segments - 1), y_segments=1)
    if ml_script1.ml_version == '1.3.4BETA':
        and_val = 'and'
    else:
        and_val = '&&'

    if script.ml_version == '1.3.4BETA': # muparser version: 1.3.2
        # Deform left side
        transform.vert_function(
            script,
            x_func='if((y>0) and (x<%s),0,x)' % (y_segments),
            y_func='if((y>0) and (x<%s),(x+1)*%s,y)' % (
                y_segments, size[1] / y_segments))
        # Deform top
        transform.vert_function(
            script,
            x_func='if((y>0) and (x>=%s),(x-%s+1)*%s,x)' % (
                y_segments, y_segments, size[0] / x_segments),
            y_func='if((y>0) and (x>=%s),%s,y)' % (y_segments, size[1]))
        # Deform right side
        transform.vert_function(
            script,
            x_func='if((y<.00001) and (x>%s),%s,x)' % (
                x_segments, size[0]),
            y_func='if((y<.00001) and (x>%s),(x-%s)*%s,y)' % (
                x_segments, x_segments, size[1] / y_segments))
        # Deform bottom
        transform.vert_function(
            script,
            x_func='if((y<.00001) and (x<=%s) and (x>0),(x)*%s,x)' % (
                x_segments, size[0] / x_segments),
            y_func='if((y<.00001) and (x<=%s) and (x>0),0,y)' % (x_segments))
    else: # muparser version: 2.2.5
        # Deform left side
        transform.vert_function(
            script,
            x_func='((y>0) && (x<{yseg}) ? 0 : x)'.format(yseg=y_segments),
            y_func='((y>0) && (x<%s) ? (x+1)*%s : y)' % (
                y_segments, size[1] / y_segments))
        # Deform top
        transform.vert_function(
            script,
            x_func='((y>0) && (x>=%s) ? (x-%s+1)*%s : x)' % (
                y_segments, y_segments, size[0] / x_segments),
            y_func='((y>0) && (x>=%s) ? %s : y)' % (y_segments, size[1]))
        # Deform right side
        transform.vert_function(
            script,
            x_func='((y<.00001) && (x>%s) ? %s : x)' % (
                x_segments, size[0]),
            y_func='((y<.00001) && (x>%s) ? (x-%s)*%s : y)' % (
                x_segments, x_segments, size[1] / y_segments))
        # Deform bottom
        transform.vert_function(
            script,
            x_func='((y<.00001) && (x<=%s) && (x>0) ? (x)*%s : x)' % (
                x_segments, size[0] / x_segments),
            y_func='((y<.00001) && (x<=%s) && (x>0) ? 0 : y)' % (x_segments))
    if center:
        transform.translate(script, [-size[0] / 2, -size[1] / 2])
    if color is not None:
        vert_color.function(script, color=color)
    return None

def half_sphere_hires():
    pass

def cube_hires(script, size=1.0, x_segments=1, y_segments=1, z_segments=1,
               simple_bottom=True, center=False, color=None):
    """Create a box with user defined number of segments in each direction.

    Grid spacing is the same as its dimensions (spacing = 1) and its
    thickness is one. Intended to be used for e.g. deforming using functions
    or a height map (lithopanes) and can be resized after creation.

    Warnings: function uses layers.join

    top_option
        0 open
        1 full
        2 simple
    bottom_option
        0 open
        1 full
        2 simple
    """
    """# Convert size to list if it isn't already
    if not isinstance(size, list):
        size = list(size)
    # If a single value was supplied use it for all 3 axes
    if len(size) == 1:
        size = [size[0], size[0], size[0]]"""
    size = util.make_list(size, 3)

    # Top
    grid(script,
         size,
         x_segments,
         y_segments)
    transform.translate(script, [0, 0, size[2]])

    # Bottom
    if simple_bottom:
        plane_hires_edges(
            script, size, x_segments, y_segments)
    else:
        layers.duplicate(script)
        transform.translate(script, [0, 0, -size[2]])
    # Rotate to correct normals
    transform.rotate(script, 'x', 180)
    transform.translate(script, [0, size[1], 0])

    # Sides
    cube_open_hires(
        script=script, size=size, x_segments=x_segments,
        y_segments=y_segments, z_segments=z_segments)

    # Join everything together
    layers.join(script)
    # Need some tolerance on merge_vert due to rounding errors
    clean.merge_vert(script, threshold=0.00002)
    if center:
        transform.translate(script, [-size[0] / 2, -size[1] / 2, -size[2] / 2])
    if color is not None:
        vert_color.function(script, color=color)
    return None

def annulus_hires(script, radius=None, radius1=None, radius2=None,
                  diameter=None, diameter1=None, diameter2=None,
                  cir_segments=48, rad_segments=1, color=None):
    """Create a cylinder with user defined number of segments

    """
    if radius is not None and diameter is None:
        if radius1 is None and diameter1 is None:
            radius1 = radius
        if radius2 is None and diameter2 is None:
            radius2 = 0
    if diameter is not None:
        if radius1 is None and diameter1 is None:
            radius1 = diameter / 2
        if radius2 is None and diameter2 is None:
            radius2 = 0
    if diameter1 is not None:
        radius1 = diameter1 / 2
    if diameter2 is not None:
        radius2 = diameter2 / 2
    if radius1 is None:
        radius1 = 1
    if radius2 is None:
        radius2 = 0
    ring = (radius1 - radius2) / rad_segments

    for i in range(0, rad_segments):
        annulus(script,
                radius1=radius1 - i * ring,
                radius2=radius1 - (i + 1) * ring,
                cir_segments=cir_segments)
    layers.join(script, merge_vert=True)
    if color is not None:
        vert_color.function(script, color=color)
    return None

def tube_hires(script, height=1.0, radius=None, radius1=None, radius2=None,
               diameter=None, diameter1=None, diameter2=None, cir_segments=32,
               rad_segments=1, height_segments=1, center=False,
               simple_bottom=False, color=None):
    """Create a cylinder with user defined number of segments

    """

    # TODO: add option to round the top of the cylinder, i.e. deform spherically
    # TODO: add warnings if values are ignored, e.g. if you specify both radius
    # and diameter.
    if radius is not None and diameter is None:
        if radius1 is None and diameter1 is None:
            radius1 = radius
        if radius2 is None and diameter2 is None:
            radius2 = 0
    if diameter is not None:
        if radius1 is None and diameter1 is None:
            radius1 = diameter / 2
        if radius2 is None and diameter2 is None:
            radius2 = 0
    if diameter1 is not None:
        radius1 = diameter1 / 2
    if diameter2 is not None:
        radius2 = diameter2 / 2
    if radius1 is None:
        radius1 = 1
    if radius2 is None:
        radius2 = 0

    # Create top
    annulus_hires(script,
                  radius1=radius1,
                  radius2=radius2,
                  cir_segments=cir_segments,
                  rad_segments=rad_segments)
    transform.translate(script, [0, 0, height])

    # Create bottom
    if simple_bottom:
        annulus(script,
                radius1=radius1,
                radius2=radius2,
                cir_segments=cir_segments)
    else:
        layers.duplicate(script)
        transform.translate(script, [0, 0, -height])
    # Rotate to correct normals
    transform.rotate(script, 'x', 180)

    # Create outer tube
    cylinder_open_hires(script, height, radius1,
                        cir_segments=cir_segments,
                        height_segments=height_segments)

    # Create inner tube
    if radius2 != 0:
        cylinder_open_hires(script, height, radius2,
                            cir_segments=cir_segments,
                            height_segments=height_segments,
                            invert_normals=True)

    # Join everything together
    layers.join(script)
    # Need some tolerance on merge_vert due to rounding errors
    clean.merge_vert(script, threshold=0.00002)
    if center:
        transform.translate(script, [0, 0, -height / 2])
    if color is not None:
        vert_color.function(script, color=color)
    return None

def triangle():
    """Create a triangle by specifying 3 points

    Under the hood: create a plane, delete one point, then move other vertices.
    """
    pass

def dna(sequence='GATTACA'):
    """Create  doublehelix function, takes spacing between helixes and
    rungs, rung diameter

    """
    pass


""" MeshLabXML utility functions """

import os
import sys
import inspect
from glob import glob

#from . import FilterScript
import meshlabxml as mlx

def is_number(num):
    """ Check if a variable is a number by trying to convert it to a float.

    """
    try:
        float(num)
        return True
    except:
        return False

def to_float(num):
    """ Convert a variable to a float """
    try:
        float(num)
        return float(num)
    except ValueError:
        return float('NaN')

def delete_all(filename):
    """ Delete all files in the current directory that match a pattern.

    Intended primarily for temp files, e.g. mlx.delete_all('TEMP3D*').

    """
    for fread in glob(filename):
        os.remove(fread)

def color_values(color):
    """Read color_names.txt and find the red, green, and blue values
        for a named color.
    """
    # Get the directory where this script file is located:
    this_dir = os.path.dirname(
        os.path.realpath(
            inspect.getsourcefile(
                lambda: 0)))
    color_name_file = os.path.join(this_dir, 'color_names.txt')
    found = False
    for line in open(color_name_file, 'r'):
        line = line.rstrip()
        if color.lower() == line.split()[0]:
            #hex_color = line.split()[1]
            red = line.split()[2]
            green = line.split()[3]
            blue = line.split()[4]
            found = True
            break
    if not found:
        print('Color name "%s" not found, using default (white)' % color)
        red = 255
        green = 255
        blue = 255
    return red, green, blue

def check_list(var, num_terms):
    """ Check if a variable is a list and is the correct length.

    If variable is not a list it will make it a list of the correct length with
    all terms identical.
    """
    if not isinstance(var, list):
        if isinstance(var, tuple):
            var = list(var)
        else:
            var = [var]
        for _ in range(1, num_terms):
            var.append(var[0])
    if len(var) != num_terms:
        print(
            '"%s" has the wrong number of terms; it needs %s. Exiting ...' %
            (var, num_terms))
        sys.exit(1)
    return var

def make_list(var, num_terms=1):
    """ Make a variable a list if it is not already

    If variable is not a list it will make it a list of the correct length with
    all terms identical.
    """
    if not isinstance(var, list):
        if isinstance(var, tuple):
            var = list(var)
        else:
            var = [var]
    #if len(var) == 1:
            for _ in range(1, num_terms):
                var.append(var[0])
    return var

def write_filter(script, filter_xml):
    """ Write filter to FilterScript object or filename

    Args:
        script (FilterScript object or filename str): the FilterScript object
            or script filename to write the filter to.
        filter_xml (str): the xml filter string

    """
    if isinstance(script, mlx.FilterScript):
        script.filters.append(filter_xml)
    elif isinstance(script, str):
        script_file = open(script, 'a')
        script_file.write(filter_xml)
        script_file.close()
    else:
        print(filter_xml)
    return None