#!/usr/bin/env python

import os, sys, time

from collections import defaultdict

import glob

import copy

import numpy as np

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

from itertools import cycle

from sklearn.metrics import roc_curve, auc

from sklearn.metrics import precision_recall_curve, average_precision_score

from sklearn.metrics import matthews_corrcoef

import rpy2.robjects as robjects

from rpy2.robjects.packages import importr

from pkg_resources import resource_string

from scipy.stats import binom

import copy

ggplot = importr('ggplot2')

importr('gridExtra')

importr('plyr')

na4com = {'A':'T', 'C':'G', 'T':'A', 'G':'C'}

def readFA(mfa, mpat='Cg', mposinpat=0, t_chr=None, t_start=None, t_end=None):

   pos_dict = defaultdict(int)

   pat3 = copy.deepcopy(mpat.upper())

   comp_pat3 = ''.join([na4com[curna] for curna in pat3][::-1])

   comp_mposinpat = len(comp_pat3)-1-mposinpat

   fadict = defaultdict();

   with open(mfa, 'r') as mr:

      cur_chr = None;

      line = mr.readline();

      while line:

         line = line.strip();

         if len(line)>0:

            if line[0]=='>': 

               if not cur_chr==None:

                  fadict[cur_chr] = ''.join(fadict[cur_chr])

               cur_chr = line[1:].split()[0]

               if t_chr in [None, cur_chr]:

                  fadict[cur_chr] = []

            else:

               if t_chr in [None, cur_chr]: 

                  fadict[cur_chr].append(line)

         line = mr.readline();

      if not cur_chr==None:

         fadict[cur_chr] = ''.join(fadict[cur_chr]) 

   fakeys = fadict.keys();

   cpgdict = defaultdict(int); cpgnum = [0, 0]

   for fak in fakeys:

       cpgdict[fak] = defaultdict()

       for i in range(len(fadict[fak])):

          if (t_start==None or i>=t_start) and (t_end==None or i<=t_end):

             if i-mposinpat>=0 and i+len(comp_pat3)-1-mposinpat<len(fadict[fak]) and ''.join(fadict[fak][i-mposinpat:(i+len(comp_pat3)-1-mposinpat+1)])==pat3:

                cpgdict[fak][('+', i)] = [1, fadict[fak][i]]; cpgnum[0] += 1

                cpgdict[fak][('-', i)] = [0, fadict[fak][i]]

             elif i-comp_mposinpat>=0 and i+len(comp_pat3)-1-comp_mposinpat<len(fadict[fak]) and ''.join(fadict[fak][i-comp_mposinpat:(i+len(comp_pat3)-1-comp_mposinpat+1)])==comp_pat3:

                cpgdict[fak][('+', i)] = [0, fadict[fak][i]]

                cpgdict[fak][('-', i)] = [1, fadict[fak][i]]; cpgnum[1] += 1

             else:

                cpgdict[fak][('+', i)] = [0, fadict[fak][i]]

                cpgdict[fak][('-', i)] = [0, fadict[fak][i]]

   print('%s%d site: %d(+) %d(-)' % (pat3, mposinpat, cpgnum[0], cpgnum[1]))

   return cpgdict

def readmodf_dict(cpgdict, modf, pred_dict, mna, t_start=None, t_end=None):

   with open(modf, 'r') as mr:

      while True:

          line = mr.readline();

          if not line: break;

          line = line.strip();

          if len(line)>0:

             lsp = line.split();

             cur_chr = lsp[0];

             cur_pos = int(lsp[1]);

             cur_strand = lsp[5];

             cur_cov = int(lsp[9]);

             cur_m_p = int(lsp[10]);

             cur_m_c = int(lsp[11]);

             if not ((t_start==None or cur_pos>=t_start) and (t_end==None or cur_pos<=t_end)):

                line = mr.readline();

                continue;

             if not (mna==lsp[3] and lsp[3]==(cpgdict[cur_chr][(cur_strand, cur_pos)][1] if cur_strand=='+' else na4com[cpgdict[cur_chr][(cur_strand, cur_pos)][1]])):

                print ('Error !! NA not equal %s == %s == %s %s' % (mna, lsp[3], cpgdict[cur_chr][(cur_strand, cur_pos)][1], modf))

             if (cur_chr, cur_pos, cur_strand) not in pred_dict: 

                pred_dict[(cur_chr, cur_pos, cur_strand)] = [cur_cov, cur_m_p, cur_m_c, lsp[3]]

             else:

                pred_dict[(cur_chr, cur_pos, cur_strand)][0] += cur_cov

                pred_dict[(cur_chr, cur_pos, cur_strand)][2] += cur_m_c

                pred_dict[(cur_chr, cur_pos, cur_strand)][1] = int(pred_dict[(cur_chr, cur_pos, cur_strand)][2]*100/pred_dict[(cur_chr, cur_pos, cur_strand)][0]) if pred_dict[(cur_chr, cur_pos, cur_strand)][0]>0 else 0

def add_from_dict(cpgdict, pred_dict, label, pred_list, mna, tp_fp_tn_fn, mpat='Cg', mposinpat=0):

   for posk in pred_dict:

             cur_chr, cur_pos, cur_strand = posk

             cur_cov, cur_m_p, cur_m_c, lsp3 = pred_dict[posk]

             iscpg = False;

             if cpgdict[cur_chr][(cur_strand, cur_pos)][0]==1:

                 iscpg = True;

                 pred_list.append((label, cur_cov, cur_m_p, cur_m_c, mpat, np.log(binom.pmf(cur_m_c, cur_cov, 0.05)) ))

                 if (lsp3==mpat[mposinpat]): pass

                 else: print ('Error not methylated pos %s %s %s' % (mna, cur_strand))

             if not iscpg:

                isclosec = False;

                for i in range(-3, 4):

                   if (cur_strand, cur_pos+i) in cpgdict[cur_chr] and cpgdict[cur_chr][(cur_strand, cur_pos+i)][0]==1:

                      isclosec = True; break;

                if lsp3==mpat[mposinpat]:

                   pred_list.append((0, cur_cov, cur_m_p, cur_m_c, mpat+'_n'+str(abs(i))+mpat[mposinpat] if isclosec else 'Other'+mpat[mposinpat], np.log(binom.pmf(cur_m_c, cur_cov, 0.05)) ))

                else:

                   pred_list.append((0, cur_cov, cur_m_p, cur_m_c, mpat+'_nb' if isclosec else 'Other', np.log(binom.pmf(cur_m_c, cur_cov, 0.05)) ))

             if pred_list[-1][0]==0:

                tp_fp_tn_fn[2] += cur_cov - cur_m_c

                tp_fp_tn_fn[1] += cur_m_c

             else:

                tp_fp_tn_fn[0] += cur_m_c

                tp_fp_tn_fn[3] += cur_cov - cur_m_c

def readmodf(cpgdict, modf, label, pred_list, mna, tp_fp_tn_fn, mpat='Cg', mposinpat=0, t_start=None, t_end=None):

   with open(modf, 'r') as mr:

      line = mr.readline();

      while line:

          line = line.strip();

          if len(line)>0:

             lsp = line.split();

             cur_chr = lsp[0];

             cur_pos = int(lsp[1]);

             cur_strand = lsp[5];

             cur_cov = int(lsp[9]);

             cur_m_p = int(lsp[10]);

             cur_m_c = int(lsp[11]); 

             if not ((t_start==None or cur_pos>=t_start) and (t_end==None or cur_pos<=t_end)): 

                line = mr.readline();

                continue;

             if not (mna==lsp[3] and lsp[3]==(cpgdict[cur_chr][(cur_strand, cur_pos)][1] if cur_strand=='+' else na4com[cpgdict[cur_chr][(cur_strand, cur_pos)][1]])):

                print ('Error !! NA not equal %s == %s == %s %s' % (mna, lsp[3], cpgdict[cur_chr][(cur_strand, cur_pos)][1], modf))

             iscpg = False;

             if cpgdict[cur_chr][(cur_strand, cur_pos)][0]==1:

                 iscpg = True;

                 pred_list.append((label, cur_cov, cur_m_p, cur_m_c, mpat, np.log(binom.pmf(cur_m_c, cur_cov, 0.05)) ))

                 if (lsp[3]==mpat[mposinpat]): pass

                 else: print ('Error not methylated pos %s %s %s' % (mna, cur_strand, modf))

             if not iscpg:

                isclosec = False;

                for i in range(-3, 4):

                   if (cur_strand, cur_pos+i) in cpgdict[cur_chr] and cpgdict[cur_chr][(cur_strand, cur_pos+i)][0]==1:

                      isclosec = True; break;

                if lsp[3]==mpat[mposinpat]:

                   pred_list.append((0, cur_cov, cur_m_p, cur_m_c, mpat+'_n'+str(abs(i))+mpat[mposinpat] if isclosec else 'Other'+mpat[mposinpat], np.log(binom.pmf(cur_m_c, cur_cov, 0.05)) ))

                else:

                   pred_list.append((0, cur_cov, cur_m_p, cur_m_c, mpat+'_nb' if isclosec else 'Other', np.log(binom.pmf(cur_m_c, cur_cov, 0.05)) ))

             if pred_list[-1][0]==0:

                tp_fp_tn_fn[2] += cur_cov - cur_m_c

                tp_fp_tn_fn[1] += cur_m_c

             else:

                tp_fp_tn_fn[0] += cur_m_c

                tp_fp_tn_fn[3] += cur_cov - cur_m_c 

          line = mr.readline();   

sssfolder = sys.argv[1];   # 

mreffile = sys.argv[2];    # 

mpat=sys.argv[3];          # Cg

mposinpat=int(sys.argv[4]);# 0 

chrofinterest = sys.argv[5];

if chrofinterest=='': chrofinterest = None;

stposofinterest = int(sys.argv[6]);

if stposofinterest<0: stposofinterest = None;

edposofinterest = int(sys.argv[7]);

if edposofinterest<0: edposofinterest = None;

basefig = sys.argv[8]

hastwoclass = 1;

sssfiles = {mpat[mposinpat]:glob.glob(os.path.join(sssfolder, 'mod_pos.*.'+mpat[mposinpat]+'.bed'))}

sssfiles[mpat[mposinpat]].extend(glob.glob(os.path.join(sssfolder, '*/mod_pos.*.'+mpat[mposinpat]+'.bed')))

sssfiles[mpat[mposinpat]].extend(glob.glob(os.path.join(sssfolder, '*/*/mod_pos.*.'+mpat[mposinpat]+'.bed')))

print(str(len(sssfiles[mpat[mposinpat]])) + " " + str(sssfolder))

## for negative;

umrfiles = []

for cur_umr_f in sys.argv[9].split(','):

   if not os.path.isdir(cur_umr_f):

       print("No prediction folder {}".format(cur_umr_f))

       sys.exit(1);

   umrfiles.extend(glob.glob(os.path.join(cur_umr_f, '*/*/mod_pos.*.'+mpat[mposinpat]+'.bed')))

   umrfiles.extend(glob.glob(os.path.join(cur_umr_f, '*/mod_pos.*.'+mpat[mposinpat]+'.bed')))

   umrfiles.extend(glob.glob(os.path.join(cur_umr_f, 'mod_pos.*.'+mpat[mposinpat]+'.bed')))

print(str(len(umrfiles)) + "  " + str(sys.argv[9].split(',')))

sys.stdout.flush()

for sa in sssfiles:

   print (sa)

   for nf in sssfiles[sa]:

      print ('\t'+nf)

cpgdict = readFA(mreffile, mpat, mposinpat, chrofinterest, stposofinterest, edposofinterest)

pred_dict = defaultdict();

for modf in umrfiles:

   readmodf_dict(cpgdict, modf, pred_dict, mpat[mposinpat], stposofinterest, edposofinterest)

baseinfo = [mpat, mpat+'_n1'+mpat[mposinpat], mpat+'_n2'+mpat[mposinpat], mpat+'_n3'+mpat[mposinpat], 'Other'+mpat[mposinpat], mpat+'_nb', 'Other']

classify_m = ['Methylation_Percentage']

classify_types = [baseinfo, [mpat]]

filename = [['all_mp','motif_mp'] ]

cov_thr = [1, 5]

mlinestyle = {1:'bo-', 3:'gx--', 5:'r*-.', 7:'cs-', 10:'md--', 15:'k+-.'}

pred_list = []; tp_fp_tn_fn = [0, 0, 0, 0]

add_from_dict(cpgdict, pred_dict, 0, pred_list, mpat[mposinpat], tp_fp_tn_fn, mpat, mposinpat)

if True:

   for na4 in sssfiles:

      for cur_f in sssfiles[na4]:

         print('%s %s' % (na4, cur_f)); sys.stdout.flush();

         readmodf(cpgdict, cur_f, hastwoclass, pred_list, na4, tp_fp_tn_fn, mpat, mposinpat, stposofinterest, edposofinterest);

   pred_list = np.array(pred_list, dtype=[('Methylation', np.uint), ('Coverage', np.uint64), ('Methylation_Percentage', np.uint64), ('Methylation_Coverage', np.uint64), ('BaseInfo', 'U20'), ('logp', np.float64)])

   if hastwoclass==1:

      cov_plot_thr = [1, 5]

      for ct_ind in range(len(classify_types)):

         ct = classify_types[ct_ind]

         cur_ct_data = pred_list[np.isin(pred_list['BaseInfo'], ct)]

         for cm_ind in range(len(classify_m)):

             print('basetype={} classify_measure={}'.format(ct, classify_m[cm_ind]))

             cm = classify_m[cm_ind]

             # 1 for roc, 2: pr;

             roc_or_pr = 2; roc_or_pr=0

             for roc_or_pr in range(1,3):

              if roc_or_pr>0: 

                mfig= plt.figure()

              if roc_or_pr==2:

                cur_fn = basefig+'/ap_plot_met_pr_'+filename[cm_ind][ct_ind]+'.png'

                xylab = ['Recall', 'Precision'];  leg_mpos = "lower left"

                for covt in cov_plot_thr:

                   precision, recall, thresholds = precision_recall_curve(cur_ct_data['Methylation'][cur_ct_data['Coverage']>=covt], cur_ct_data[cm][cur_ct_data['Coverage']>=covt])

                   ap_pr = average_precision_score(cur_ct_data['Methylation'][cur_ct_data['Coverage']>=covt], cur_ct_data[cm][cur_ct_data['Coverage']>=covt])

                   plt.plot(recall, precision, mlinestyle[covt], lw=2, label='Coverage>=%d (AP=%0.3f)' % (covt, ap_pr))

                   print('\t\t %s %d ap=%.5f' % (cur_fn, covt, ap_pr))

              elif roc_or_pr==1:

                xylab = ['False Positive Rate', 'True Positive Rate']; leg_mpos = "lower right"

                cur_fn = basefig+'/roc_plot_met_roc_'+filename[cm_ind][ct_ind]+'.png'

                prev = 0; prev_ind = -1

                for covt in cov_plot_thr:

                   fpr, tpr, mthr = roc_curve(cur_ct_data['Methylation'][cur_ct_data['Coverage']>=covt], cur_ct_data[cm][cur_ct_data['Coverage']>=covt])

                   #print(','.join([str(np.round(t1, 5)) for t1 in mthr]))

                   roc_auc = auc(fpr, tpr)

                   if (not np.isnan(roc_auc)) and (abs(roc_auc - prev)>=0.02 or (covt>10 and abs(roc_auc - prev)>=0.005) or (cov_plot_thr.index(covt)-prev_ind>1 and abs(roc_auc - prev)>=0.005)):

                      plt.plot(fpr, tpr, mlinestyle[covt], lw=2, label='Coverage>=%d (AUC=%0.3f)' % (covt, roc_auc)) 

                      prev = roc_auc; prev_ind = cov_plot_thr.index(covt)

                   if not np.isnan(roc_auc):

                      print ('\t\t %s %d %.7f' % (cur_fn, covt, roc_auc))

                plt.plot([0, 1], [0, 1])

              if roc_or_pr>0:

                plt.xlim([0.0, 1.0]);              plt.ylim([0.0, 1.0])

                plt.xlabel(xylab[0]);              plt.ylabel(xylab[1])

                plt.legend(loc=leg_mpos)

                mfig.savefig(cur_fn, dpi=300);              plt.close(mfig)