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
import subprocess

def scale(X):
  m25 = np.percentile(X[:,0], 25)
  m75 = np.percentile(X[:,0], 75)
  s50 = np.median(X[:,2])
  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")
  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))
  f.close()

  print "s", s
  ret = subprocess.call("ulimit -v 32000000; ./realign <%stmpb-%s.in >%stmpb-%s.out" % (base_dir, s, base_dir, s), shell=True)

  if ret != 47:
    print "fail", s
    return

  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()

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[2:]:
    f = open(fn)
    ref = f.readline()
    if len(ref) > 30000:
      print "out", len(ref)
      continue
    refs.append(ref.strip())
    names.append(fn)
    X = []
    Y = []
    Y2 = []
    print "\rfn", fn,
    sys.stdout.flush()
    for l in f:
      its = l.strip().split()
      mean = float(its[0])
      std = float(its[1])
      length = float(its[2])
      X.append([mean, mean*mean, std, length])
      Y.append(4)
      Y2.append(4)
    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 ("done", sum(len(x) for x in refs))
  sys.stdout.flush()

  ntwk = BatchNet(n_dims, n_classes, only_test=True)
  ntwk.load(sys.argv[1])
  base_dir = str(datetime.datetime.now())
  base_dir = base_dir.replace(' ', '_')

  os.mkdir(base_dir)
  base_dir += "/"

  p = Pool(8)
  p.map(realign, range(len(data_x)))
