#Usage: python analyseEA.py (your miRNA list file) (The miRNA weighting score)
#The miRNA list of interest needs to be converted into precursor level


# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import operator,io,math
import copy
from scipy import special,stats
from pandas.compat import StringIO


def gammaln(c_bottom,c_top):
    c_bottom+=1
    c_top+=1
    return special.gammaln(c_bottom)-special.gammaln(c_bottom-c_top+1)-special.gammaln(c_top)


def hyper_gamma(N,M,n,m):
    return np.exp(gammaln(n,m)+gammaln(N-n,M-m)-gammaln(N,M))



def hyper_gamma_overlep(N,M,n,m):
    if n > M:
        tmp=math.floor(M-m)
    else:
        tmp=math.floor(n-m)
    k=tmp+m+1
    return sum([hyper_gamma(N, M, n, i) for i in np.arange(m, k)])



def set_handler1(file):
    d = {}
    for line in file:
        a = line.strip().split('\t', 2)
        curr_set=a[2].split('\t')
        if len(curr_set)>=2:
            d[a[0]] = d.get(a[0], {'name': a[0], 'all':set(),'items':{}})
            d[a[0]]['items'][a[1]] = set(curr_set)
    return d


def set_handler2(set_d, l=1, u=np.inf):
    num=0
    sets_d = copy.deepcopy(set_d)
    for i,j in sets_d.items():
        for x, y in j['items'].items():
            if l<=len(y)<=u:
                num+=1
                sets_d[i]['all'] = sets_d[i]['all'].union(y)
            else:
                print x
                del j['items'][x]
    return sets_d,num


all_sets = set_handler1(io.open(r'mirset.txt','r',encoding='utf-8'))
all_sets,sets_n=set_handler2(all_sets)
print sets_n
print len(all_sets)


def normalized(score_list):
    norm_min=np.min(score_list)
    norm_max=np.max(score_list)
    norm_fm=norm_max-norm_min
    return [(x-norm_min)*1.0/norm_fm for x in score_list]


def secondNorm(score_list):
    norm_mean=np.mean(score_list)
    return [i-norm_mean+1 for i in score_list]


def fdr_pre(list, n):
    return [[i[n], list.index(i)] for i in list]


def fdr(list1, sets_n):
    list1.sort()
    list2 = []
    n = 1
    for i, j in list1:
        list2.append(i*sets_n/n)
        n = n + 1
    map(list.append, list1, list2)
    return list1

def self_format(x):
    return round(x,2) if x>0.01 else format(x,".2e")



def rich_out(input_lst,score_file,factor=1,udef=None):
    result = []
    if 'mic' not in score_file:
        if 'udef' in score_file:
            scoredf = pd.read_csv(StringIO(udef),header=None)
            scoredf.columns=['mir','score']
            print scoredf.shape
            print scoredf
            scoredf.loc[:, 'score'] = normalized(scoredf['score'])
        else:
            scoredf = pd.read_table( score_file, header=0)
            scoredf.loc[:, 'score'] = normalized(scoredf['score'])
    else:
        scoredf = pd.read_table(score_file, header=0)
        back=pd.read_table(r'Allset_background.txt',header = 0)
    back_id=set(back.iloc[:,0])
    scoredf=scoredf[scoredf['mir'].isin(back_id)]
    scoredf.loc[:,'score']=secondNorm(scoredf['score'])
    print type(scoredf['score'])
    scoredf.loc[:,'score_fac']=np.power(scoredf['score'],float(factor))
    lost_mirs=back[back.iloc[:,0].isin(back_id-set(list(scoredf.loc[:,'mir'])))]
    lost_mirs.reset_index(inplace=True)
    tmp=[]
    [tmp.append([lost_mirs.iloc[i,0],1,1]) for i in range(len(lost_mirs))]
    scoredf=pd.concat([scoredf[['mir','score','score_fac']],pd.DataFrame(tmp,columns=['mir','score','score_fac'])],axis=0,ignore_index=True)
    input_set=set(input_lst)
    print len(input_set)
    a = input_set & back_id 
    input_score=scoredf[scoredf.mir.isin(a)]
    c = back_id
    N_o = len(c)
    M_o = len(a) 
    N=sum(scoredf[scoredf['mir'].isin(c)].loc[:,'score_fac'])
    M=sum(scoredf[scoredf['mir'].isin(a)].loc[:,'score_fac'])
    for all_set in all_sets:
        all_set = all_sets[all_set]
        for i,j in all_set['items'].items():
            b = a & j#输入与功能set
            n_o = len(j & c)
            m_o = len(b)
            n=sum(scoredf[scoredf['mir'].isin(j & c)].loc[:,'score_fac'])
            m=sum(scoredf[scoredf['mir'].isin(b)].loc[:,'score_fac'])
            hited =input_score[input_score.mir.isin(b)]
            hited_str=''
            for h in range(len(hited)):
                if h!=len(hited)-1:
                    hited_str+=hited.iloc[h,0]+'('+str(round(hited.iloc[h,-1],3))+'),'
                else:
                    hited_str+=hited.iloc[h,0]+'('+str(round(hited.iloc[h,-1],3))+')'
            missed =scoredf[scoredf.mir.isin(j - a)]
            miss_str=''
            for h in range(len(missed)):
                if h!=len(missed)-1:
                    miss_str+=missed.iloc[h,0]+'('+str(round(missed.iloc[h,-1],3))+'),'
                else:
                    miss_str+=missed.iloc[h,0]+'('+str(round(missed.iloc[h,-1],3))+')'
            if m > 0:
                p_v = hyper_gamma_overlep (N, M, n, m)
                p_v_o=stats.fisher_exact([[m_o,M_o-m_o],[n_o-m_o,N_o-M_o-n_o+m_o]],alternative='greater')[1]
                result.append([all_set['name'], i,str(m_o)+'/'+str(n_o),str(self_format(m))+'/'+str(self_format(n)),
                               float(m)*N/(n*M),p_v,p_v,p_v,hited_str,miss_str,p_v_o])
    lenre=len(result)
    for i in fdr(fdr_pre(result, 5), sets_n):
        result[i[1]][6]=i[2]
        result[i[1]][7]=result[i[1]][7]*lenre
    for i in fdr(fdr_pre(result, 10), sets_n):
        result[i[1]].append(i[2])
        result[i[1]].append(result[i[1]][10]*lenre)
    result.sort()
    dic_res={}
    dic_res['raw_sets']=len(input_set)
    dic_res['inter_sets'] = len(a)
    dic_res['result']=result
    pd.DataFrame(dic_res['result']).to_csv('static/test.csv',index=False,header=False)
    return dic_res

if (__name__ == "__main__"):
    try:
        miR_path = sys.argv[1]
        score_path = sys.argv[2]
        miR_lst=pd.read_table(miR_path,header=None).iloc[:,0]
        LRMODEL_FEATURE = LRMODEL_FEATURES[species]
        LRMODEL_TRIPLETS = LRMODEL_FEATURE[2:7]
        LRMODEL = LRMODELS[species]
        res=pd.DataFrame(rich_out(miR_lst,score_path,3),columns=['class','term','Count','Percent','OR','WEA_P','WEA_FDR','WEA_bonf','Hit','Miss','NEA_P','NEA_FDR','NEA_bonf'])
        res.to_csv("wTAM_result.txt",index=False)
    except IndexError, e:
        print '''python analyseEA.py (your miRNA list file) (The miRNA weighting score)
