#!/usr/bin/env python

import os, sys, string, time
from collections import defaultdict;
import multiprocessing



def read_genome(mfafile):
   ref_genome = defaultdict();
   with open(mfafile, 'r') as mr:
      cur_chr = None;
      while True:
         line = mr.readline();
         if not line: break;
         line = line.strip();
         if len(line)==0: continue;
         if line[0]=='>':
            if not cur_chr==None:
               ref_genome[cur_chr] = ''.join(seqlist);
            cur_chr = line[1:].split()[0];
            seqlist = []
         else:
            seqlist.append(line.upper())
      ref_genome[cur_chr] = ''.join(seqlist);
   print("Total chr: {}".format(len(ref_genome))); sys.stdout.flush()
   return ref_genome
#ref_genome = read_genome(ref_fa);

def handle_motif_pos(run_Q):
   handli = 0;
   while not run_Q.empty():
      try:
         rgkey, ref_genome, res_folder, na_bp, curna, curmotif = run_Q.get(block=False)
         #print(rgkey, ref_genome, res_folder, na_bp, curna, curmotif); continue
      except:
         break;

      #curna_dict = defaultdict();
      #curmotif_dict = defaultdict();
      nafile = '%sna_%s_%s.bed' % (res_folder, rgkey, curna)
      motiffile = '%smotif_%s_%s.bed' % (res_folder, rgkey, curna)
      mw_na = open(nafile, 'w')
      mw_motif = open(motiffile, 'w')

      batchsize = 500000
      print("get motif for {}={}".format(rgkey, len(ref_genome))); sys.stdout.flush()
      cur_hi = 0; start_time = time.time();
      for na_ind in range(len(ref_genome)):
          cur_hi += 1
          if cur_hi % batchsize == 0:
             print('\t time consuming ({})= {} {}'.format( rgkey, cur_hi, time.time() - start_time ) )
             sys.stdout.flush()
             start_time = time.time()

          if (ref_genome[na_ind]==curna or na_bp[ref_genome[na_ind]]==curna):
             #curna_dict[(rgkey, na_ind, '+' if ref_genome[rgkey][na_ind]==curna else '-')] = True;
             mw_na.write('%s\t%s\t%s\n' % (rgkey, na_ind, '+' if ref_genome[na_ind]==curna else '-'))
          if ref_genome[na_ind]==curna and (not curmotif==None):
             for cur_mot in curmotif:
                 is_mot = True; mot_ind = 0;
                 for cur_sub_r_ind in range(na_ind - curmotif[cur_mot], na_ind + len(cur_mot) - curmotif[cur_mot] ):
                     if cur_sub_r_ind<0 or cur_sub_r_ind>len(ref_genome)-1: 
                        is_mot = False; break;
                     if not ref_genome[cur_sub_r_ind] == cur_mot[mot_ind]: 
                        is_mot = False; break;
                     mot_ind += 1
                 if is_mot:
                    mw_motif.write('%s\t%s\t%s\n' % (rgkey, na_ind, '+'))
                    mw_motif.write('%s\t%s\t%s\n' % (rgkey, na_ind+1, '-'))
                    break;
      mw_na.close(); mw_motif.close()

                   

ref_fa = 'ref/hg38.fa'
ref_fa = sys.argv[1]
res_folder = 'genome.motif/C/'
res_folder = sys.argv[2]+'/'
if not os.path.isdir(res_folder):
   os.system('mkdir -p '+res_folder)

curna='C'
curmotif={'CG':0}
curna = sys.argv[3];
curmotif = {sys.argv[4]:int(sys.argv[5])}

if len(sys.argv)>6:
   chrkeys = ["chr%s" % cid for cid in sys.argv[6].split(',')]
else:
   chrkeys = []
   for i in range(1, 23):
      chrkeys.append("chr%d" % i)
   chrkeys.append("chrX")
   chrkeys.append("chrY")
   chrkeys.append("chrM")

chrkeys = set(chrkeys)


na_bp = {"A":"T", \
         "C":"G", \
         "G":"C", \
         "T":"A", \
         "a":"t", \
         "c":"g", \
         "g":"c", \
         "t":"a", \
         "N":"N", \
         "n":"n" \
         }


ref_genome = read_genome(ref_fa);

##############################
pmanager = multiprocessing.Manager();
run_Q = pmanager.Queue();
for curk in chrkeys:
   run_Q.put((curk, ref_genome[curk], res_folder, na_bp, curna, curmotif))

mhandlers = [];
share_var = (run_Q, )
m_thread_num = len(chrkeys);
for i in range(m_thread_num):
   p = multiprocessing.Process(target=handle_motif_pos, args=share_var)
   p.start();
   mhandlers.append(p);
while any(p.is_alive() for p in mhandlers):
   try:
      time.sleep(1);
   except:
      time.sleep(1);
      continue;