"""
    作者：徐诗芬
    内容：1.先做一个右尾区的poisson概率检验，每读取两行为一组，第一行的数值为观测值obs，第二行的数值为期望值exp，
         以exp为mu,去检验k≥obs的概率p = 1-poisson.cdf(k = obs, mu = exp)+poisson.pmf(k = obs, mu = exp),且每一行有两个数值，
         一个是parallel，一个是广义的convergence，分别进行检验；
         2.做完检验后再进行多重检验校正，一般使用Benjamini-Hochberg FDR (BH法)，即qvalue = pvalue*length/rank(p)，
         并筛选p<=0.05的基因

    日期：2022.2.23
"""
from scipy.stats import poisson
import re
import sys

def usage():
    print('Usage: python result.py [inputfile] [outfile1] [outfile2]')

def gp(file):
    with open(file, 'rt') as f:
        f = f.readlines()[1:]   # 跳过第一行并生成列表
        gp = map(lambda x: x.strip().split('\t'), f)  # 这里的map是一个迭代器，批量读取文件里的元素并进行切割
        return gp       # 迭代器放在自定义函数里，不要放在main里面，可以多次使用

def main():
    infile = open(sys.argv[1], 'rt')     # 'gene.log'
    outfile1 = sys.argv[2]               # 'gene.pvalue.txt'
    outfile2 = open(sys.argv[3], 'wt')   # 'gene.qvalue.txt'
    # dat_dict = {}
    # 1. 进行poisson test
    with open(outfile1, 'wt') as of1:
        of1.write('gene'+'\t'+'pvalue_con' + '\t'+'pvalue_par' + '\n')
        for line in infile:
            line_obs = line.strip().split('\t')
            name = re.search(r'OG[0-9]*', line_obs[0]).group()
            obs_p = eval(line_obs[2])
            obs_c = eval(line_obs[3])
            line = infile.readline()
            line_exp = line.strip().split('\t')
            exp_p = eval(line_exp[2])
            exp_c = eval(line_exp[3])
            pvalue_c = 1 - poisson.cdf(obs_c, exp_c) + poisson.pmf(obs_c, exp_c)
            pvalue_p = 1 - poisson.cdf(obs_p, exp_p) + poisson.pmf(obs_p, exp_p)
            new_line = '{}\t{}\t{}'.format(name, pvalue_c, pvalue_p)
            # 把广义的convergence检验放在前面，先关注它的结果
            of1.write(new_line + '\n')
        infile.close()

    # 2. multiple test correction
    outfile2.write(
        'gene' + '\t' + 'pvalue_con' + '\t' + 'qvalue_con' + '\t' + 'pvalue_par' + '\t' + 'qvalue_par' + '\n')
    p_con_list = []
    p_par_list = []
    for i in gp(outfile1):
        p_con_list.append(eval(i[1]))
        p_par_list.append(eval(i[2]))
    length = len(p_con_list)
    p_con_list.sort(reverse=False)  # 对原列表进行排序
    p_par_list.sort(reverse=False)
    flag = 0
    for i in gp(outfile1):
        # 计算BH-FDR校正后的adj.p
        qvalue_con = eval(i[1]) * length / (p_con_list.index(eval(i[1])) + 1)
        qvalue_par = eval(i[2]) * length / (p_par_list.index(eval(i[2])) + 1)
        if qvalue_con <= 0.05 or qvalue_par <= 0.05:
            newline = '{}\t{}\t{}\t{}\t{}'.format(i[0], i[1], qvalue_con, i[2], qvalue_par)
            outfile2.write(newline + '\n')
            flag += 1
    outfile2.close()

try:
    main()
except IndexError:
    usage()