from qrnn import BatchNet
import pickle
import sys
import numpy as np
import datetime
from collections import defaultdict
import os
from sklearn.metrics import confusion_matrix
import theano as th
from multiprocessing import Pool

def scale(X):
  m25 = np.percentile(X[:,0], 25)
  m75 = np.percentile(X[:,0], 75)
  s50 = np.median(X[:,2])
#  me25 = 0.07499809
#  me75 = 0.26622871
  me25 = -0.3
  me75= 0.3
  se50 = 0.6103758
  ret = np.array(X)
  scale = (me75 - me25) / (m75 - m25)
  m25 *= scale
  shift = me25 - m25
  ret[:,0] = X[:,0] * scale + shift
  ret[:,1] = ret[:,0]**2
  
  sscale = se50 / s50

  ret[:,2] = X[:,2] * sscale
  return ret

def print_stats(o):
  stats = defaultdict(int)
  for x in o:
    stats[x] += 1
  print (stats)

def flatten2(x):
  return x.reshape((x.shape[0] * x.shape[1], -1))

def realign(s):
  ps = s 
  o1, o2 = ntwk.tester([scale(data_x[ps])]) 
  o1m = (np.argmax(o1[0], 1))
  o2m = (np.argmax(o2[0], 1))
  alph = "ACGTN"
#      fo = open(base_dir+"output%s.fasta" % ps, "w")
#      print (">xx", file=fo)
#      for a, b in zip(o1m, o2m):
#        if a < 4: 
#          fo.write(alph[a])
#        if b < 4: 
#         fo.write(alph[b])
#      fo.close()      
  f = open(base_dir+"tmpb-%s.in" % s, "w")
  lc = 0
  print >>f, refs[ps]
  for a, b in zip(o1[0], o2[0]):
    print >>f, " ".join(map(str, a))
    print >>f, " ".join(map(str, b))
    lc += 1
  f.close()

  print "s", s, datetime.datetime.now(), len(refs[ps]), lc, len(data_x[ps])
  if os.system("./realign %stmpb-%s.in <%stmpb-%s.in >%stmpb-%s.out" % (base_dir, s, base_dir, s, base_dir, s)) != 0:
    print "watwat", s
    sys.exit()

  f = open(base_dir+"tmpb-%s.out" % s)
  for i, l in enumerate(f):
    data_y[ps][i] = mapping[l[0]]
    data_y2[ps][i] = mapping[l[1]]

  fo = open(names[s] + "r", "w")
  print >>fo, refs[s]
  for x, y, y2 in zip(data_x[s], data_y[s], data_y2[s]):
    print >>fo, " ".join(map(str, x)), "%c%c" % (alph[y], alph[y2])
  fo.close()
  return data_y[ps], data_y2[ps]

if __name__ == '__main__':
  chars = "ACGT"
  mapping = {"A": 0, "C": 1, "G": 2, "T": 3, "N": 4}

  n_dims = 4
  n_classes = 5

  data_x = []
  data_y = []
  data_y2 = []
  refs = []
  names = []

  for fn in sys.argv[1:]:
    if fn.endswith("pkl"):
      continue
    f = open(fn)
    ref = f.readline()
    if len(ref) > 30000:
      print "out", len(ref)
      continue
    X = []
    Y = []
    Y2 = []
    good = True
    for l in f:
      its = l.strip().split()
      X.append(list(map(float, its[:4])))
      try:
        Y.append(mapping[its[4][0]])
        Y2.append(mapping[its[4][1]])
      except:
        print "wat", fn
        good = False
        break
    if good:
      if len(X) > 2000:
        print "\rfn", fn, len(X), len(Y), len(Y2), len(ref), len(refs),
        sys.stdout.flush()
        refs.append(ref.strip())
        names.append(fn)
        data_x.append(np.array(X, dtype=th.config.floatX))
        data_y.append(np.array(Y, dtype=np.int32))
        data_y2.append(np.array(Y2, dtype=np.int32))

  print
  print ("done", sum(len(x) for x in refs), sum(len(x) for x in data_x))
  sys.stdout.flush()

  ntwk = BatchNet(n_dims, n_classes)
  if sys.argv[1].endswith("pkl"):
    ntwk.load(sys.argv[1])

  print ("net rdy")
  sys.stdout.flush()

  s_arr = []
  p_arr = []
  subseq_size = 2000
  for s in range(len(data_x)):
    s_arr += [s]
    p_arr += [len(data_x[s]) - subseq_size]

  sum_p = sum(p_arr)
  for i in range(len(p_arr)):
    p_arr[i] = 1.*p_arr[i] / sum_p

  base_dir = str(datetime.datetime.now())
  base_dir = base_dir.replace(' ', '_')

  os.mkdir(base_dir)
  base_dir += "/"
  batch_size = 8
  n_batches = len(data_x) / batch_size + 1
  print len(data_x), batch_size, n_batches, datetime.datetime.now()

  for epoch in range(1000):
    taken_gc = []
    out_gc = []
    tc = 0
    tcs = 0
    tcb = 0
    tc2 = 0
    tc3 = 0
    ccb = 0
    ccs = 0
    total_size = 0
    o1mm = []
    y1mm = []
    o2mm = []
    y2mm = []
    for s in range(n_batches):
      bx = []
      by = []
      by2 = []
      if s % 10 == 0:
        s2 = np.random.choice(s_arr)
        x = data_x[s2]
        y = data_y[s2]
        y2 = data_y2[s2]
        total_size += len(x)
        bx.append(scale(x))
        by.append(y)
        by2.append(y2)
      else:
        for b in range(batch_size):
          s2 = np.random.choice(s_arr, p=p_arr)
          r = np.random.randint(0, data_x[s2].shape[0] - subseq_size)
          x = data_x[s2][r:r+subseq_size]
          total_size += len(x)
          y = data_y[s2][r:r+subseq_size]
          y2 = data_y2[s2][r:r+subseq_size]
          bx.append(scale(x))
          by.append(y)
          by2.append(y2)

      lr = 8e-3
#      lr = 1e-3
  #    if epoch >= 3:
  #      lr = 2e-1
      if epoch >= 970:
        lr = 1e-3
      by = np.array(by)
      by2 = np.array(by2)
      cost, o1, o2 = ntwk.trainer(bx, by, by2)
      tc += cost
      if s % 20 == 0:
        tcb += cost
        ccb += 1
      else:
        tcs += cost
        ccs += 1

      o1m = (np.argmax(flatten2(o1), 1))
      o2m = (np.argmax(flatten2(o2), 1))

      o1mm += list(o1m)
      o2mm += list(o2m)
      y1mm += list(flatten2(by))
      y2mm += list(flatten2(by2))

      tc2 += np.sum(np.equal(o1m, by.flatten()))
      tc3 += np.sum(np.equal(o2m, by2.flatten()))

      sys.stdout.write('\r%d %f %f %f' % (s, tc / (s+1), tcs / max(1, ccs), tcb / max(1, ccb)))
      sys.stdout.flush()

    print
    conf1 = confusion_matrix(y1mm, o1mm)
    conf2 = confusion_matrix(y2mm, o2mm)
    good = conf1[0,0] + conf1[1,1] + conf1[2,2] + conf1[3,3]
    good += conf2[0,0] + conf2[1,1] + conf2[2,2] + conf2[3,3]
    bad = np.sum(conf1) + np.sum(conf2) - good - conf1[4,4] - conf2[4,4]

    print epoch, tc / n_batches, 1.*tc2 / total_size, 1.*tc3 / total_size, 1.*good / (good + bad), datetime.datetime.now()
    print_stats(o1mm)
    print_stats(o2mm)
    print conf1
    print conf2
  #  print "out", np.min(out_gc), np.median(out_gc), np.max(out_gc), len(out_gc)
    sys.stdout.flush()

    if epoch % 20 == 19:
      ntwk.save(base_dir+"dumpx-%d.pkl" % epoch)
    ntwk.save(base_dir+"latest.pkl")
