import os
import argparse
import pandas as pd
import numpy as np

HELP="""
"""

def get_p_d_cluster_id(cluster_ids, strand):
    temp_dict = {int(i.split('_')[-1]):i for i in cluster_ids}
    keys = sorted(temp_dict)
    clusters = [temp_dict[i]for i in keys] if strand == '+' else [temp_dict[i]for i in keys[::-1]]
    return clusters

parser = argparse.ArgumentParser(description=HELP, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("-c", "--control", type=str, nargs="+", required=True,
                    help="a list of control files generated by <6.clustering_cleavage_sites.py>, al least 2.")
parser.add_argument("-t", "--treatment", type=str, nargs="+", required=True,
                    help="a list of treatment files generated by <6.clustering_cleavage_sites.py>, at least 2.")
parser.add_argument("-o", "--out_path", type=str, default="./",
                    help="The output path.")
parser.add_argument("-d", "--distance", type=int, default=0,
                    help="The distance between clusters in each sample to merge as a single site.")
parser.add_argument("-p", "--p_value", type=float, default=0.05,
                    help="The threshold of significance in DEXSeq, should be within (0, 1)")
parser.add_argument("-padj", "--p_adjust", type=float, default=1,
                    help="The threshold of significance in DEXSeq, should be within (0, 1)")
parser.add_argument("-fc", "--log2fc", type=float, default=0,
                    help="The threshold of fold change (log2 transformed).")
args = parser.parse_args()

#0. argument check
if not os.path.exists(args.out_path):
    os.makedirs(args.out_path)
if not 0 < args.p_value <= 1:
    raise ValueError("Invalid argument paramenter: -p should be in (0, 1].")
if not 0 < args.p_adjust <= 1:
    raise ValueError("Invalid argument paramenter: -padj should be in (0, 1].")
if not args.distance >= 0:
    raise ValueError("Invalid argument paramenter: -d should at least 0.")
if not args.log2fc >= 0:
    raise ValueError("Invalid argument paramenter: -fc should at least 0.")
if not (len(args.control) > 1 or len(args.treatment) > 1):
    raise ValueError("Invalid argument paramenter: The number of control or treatment file should be both at least 2.")

#1. generate annotation <*.gff> file and count data <*.count> file
k = 0
for i in args.control:
    prefix = os.path.split(i)[-1].split('.')[0]
    out_name = os.path.join( args.out_path, 'control_%d.add_label.bed'%(k) )
    cmd = r"""awk 'BEGIN{OFS="\t"}{print $1, $2, $3, $10, $5, $6, "control_%d"}' %s >%s """%(k, i, out_name)
    print(cmd)
    os.system(cmd)
    k += 1
for i in args.treatment:
    prefix = os.path.split(i)[-1].split('.')[0]
    out_name = os.path.join( args.out_path, 'treatment_%d.add_label.bed'%(k) )
    cmd = r"""awk 'BEGIN{OFS="\t"}{print $1, $2, $3, $10, $5, $6, "treatment_%d"}' %s >%s """%(k, i, out_name)
    print(cmd)
    os.system(cmd)
    k += 1
#1.1 merge robust clusters
cmd = r"""less -S %s/*.add_label.bed |bedtools sort |bedtools cluster -s -d %d """%(args.out_path, args.distance)
samples = [ i.split('.')[0] for i in os.listdir(args.out_path) if i.endswith('.add_label.bed') ]
stat = {}
for line in os.popen(cmd).read().strip().split('\n'):
    chrom, start, end, gene_id, count, strand, label, cluster = line.strip().split('\t')
    gene_id = gene_id+"|%s"%strand
    start, end = int(start), int(end)
    count = int(count)
    cluster = "mCluster_%s"%(cluster)

    # gene_id
    if not gene_id in stat:
        stat[gene_id] = {'chrom':chrom, 'itv':[start, end], 'strand':strand, 'exons':{} }
    if end > stat[gene_id]['itv'][-1]:
        stat[gene_id]['itv'][-1] = end
    #cluster_id
    if not cluster in stat[gene_id]['exons']:
        stat[gene_id]['exons'][cluster] = {'itv':[start, end], 'samples':{} }
    if end > stat[gene_id]['exons'][cluster]['itv'][-1]:
        stat[gene_id]['exons'][cluster]['itv'][-1] = end
    #sample_id
    if not label in stat[gene_id]['exons'][cluster]['samples']:
        stat[gene_id]['exons'][cluster]['samples'][label] = 0
    stat[gene_id]['exons'][cluster]['samples'][label] += count
#1.2 generate files
gff_name = os.path.join(args.out_path, 'sample_merge_robust_cluster_for_DEXSeq.gff')
gff = open(gff_name, 'w')
count_files = {i:open(os.path.join(args.out_path, i+'.count'), 'w') for i in samples}
#gene_id
for gene_id in stat:
    chrom, start, end, strand = stat[gene_id]['chrom'], stat[gene_id]['itv'][0], stat[gene_id]['itv'][-1], stat[gene_id]['strand']
    outline = [chrom, 'PASfinder', 'aggregate_gene', str(start+1), str(end), '.', strand, '.', 'gene_id "%s"'%(gene_id)]
    gff.write('\t'.join(outline) + '\n')
#cluster_id
    exons = stat[gene_id]['exons']
    trans_id = gene_id+"|%d"%(len(exons))
    for n, cluster_id in enumerate(exons):
        start, end = exons[cluster_id]['itv']
        outline = [chrom, 'PASfinder', 'exonic_part', str(start+1), str(end), '.', strand, '.', 'transcripts "%s"; exonic_part_number "%s"; gene_id "%s"'%(trans_id, cluster_id, gene_id) ]
        gff.write('\t'.join(outline) + '\n')
#sample_id
        tmp = exons[cluster_id]['samples']
        for key in samples:
            outline = ['%s:%s'%(gene_id, cluster_id), str(tmp.get(key, 0) )]
            count_files[key].write('\t'.join(outline) + '\n')
gff.close()
for key in count_files:
    count_files[key].close()

#2. run DEXSeq.r
cmd = 'Rscript DEXSeq.r ./%s/'%(args.out_path)
#os.system(cmd)

#3. select APA events
stat = {}
for n, line in enumerate(open(os.path.join(args.out_path, 'DEXSeq_treatment_vs_control.txt'), 'r')):
    if n == 0:
        title = line.strip().split('\t')
        title = [title[1], title[2], title[5]] #pval, padj, log2fc
        continue
    s = line.strip().split("\t")
    s[-1] = 0 if s[-1] == "NA" else s[-1]
    trans_id, cluster_id = s[0].split(':E')
    if not trans_id in stat:
        stat[trans_id] = {}
    stat[trans_id][cluster_id] = dict(zip(title, list(map(float, [s[2], s[3], s[6]]))))

out = open(os.path.join(args.out_path, 'DEXSeq_treatment_vs_control.diff'), 'w')
for trans_id in stat:
    if len(stat[trans_id]) == 1:
        continue

    strand = trans_id[-1]
    trans = pd.DataFrame(stat[trans_id])
    trans_filter = trans.loc[:, trans.loc[title[0], :] < args.p_value ]                          #pvalue
    if trans_filter.shape[1] > 0:
        trans_filter = trans_filter.loc[:, trans_filter.loc[title[0], :] < args.p_adjust ]       #padjust
    if trans_filter.shape[1] > 0:
        trans_filter = trans_filter.loc[:, abs(trans_filter.loc[title[0], :]) > args.log2fc ]    #log2foldchange

    if trans_filter.shape[1] < 2:
        change = 'no'
        cluster_ids = get_p_d_cluster_id(trans.columns.tolist(), strand)
        trans_filter_top2 = trans.loc[:, cluster_ids]
    else:
        trans_filter = trans_filter.sort_values(axis=1, by=title[0], ascending=True)
        trans_filter_top2 = trans_filter.iloc[:, [0,1]]
        cluster_ids = get_p_d_cluster_id(trans_filter_top2.columns.tolist(), strand)
        trans_filter_top2 = trans.loc[:, cluster_ids]
        if np.alltrue(trans_filter_top2.loc[title[2]]<0) or np.alltrue(trans_filter_top2.loc[title[2]]>0):
            change = 'same_change_for_both_proximal_and_distal_site' 
        elif trans_filter_top2.loc[title[2], cluster_ids[0]] < 0:
            change = 'short'
        else:
            change = 'long'
    outline = [trans_id, change,               
               ','.join(trans_filter_top2.columns.tolist()),
               ','.join(map(str, trans_filter_top2.loc[title[2], :].tolist())),
               ','.join(map(str, trans_filter_top2.loc[title[0], :].tolist())),
               ','.join(map(str, trans_filter_top2.loc[title[1], :].tolist())) ] 
    out.write('\t'.join(outline) + '\n' )
out.close()
