import msgpack
import json
import pickle
import os.path
from Queue import PriorityQueue
import re
import doench_score
import azimuth.model_comparison
import numpy as np
import pandas as pd
import csv
from intervaltree import IntervalTree
from multiprocessing import Process
import os
import time

start_time = time.time()

#Reverse complements a given string
def revcom(s):
    basecomp = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A','U':'A', 'N':'N'}
    letters = list(s[::-1])
    letters = [basecomp[base] for base in letters]
    return ''.join(letters)

class GuideRNA():
  """Holder of gRNA information"""
  def __init__(self, selected, start, seq, PAM, seq_before, seq_after, chrom, cut_pos, score, exon_ranking, ensembl_gene, gene_name, functional_domain, has_exome_repeat, off_target_score):
    self.start = start
    self.seq = seq
    self.PAM = PAM
    self.seq_before = seq_before # 10bp before the sgRNA
    self.seq_after = seq_after # 10bp after the sgRNA
    self.chrom = chrom
    self.cut_pos = cut_pos
    self.score = score
    self.exon_ranking = exon_ranking
    self.ensembl_gene = ensembl_gene
    self.gene_name = gene_name
    self.selected = selected
    self.functional_domain = functional_domain
    if functional_domain:
      self.has_functional_domain = True
    else:
      self.has_functional_domain = False
    self.has_exome_repeat = has_exome_repeat
    self.off_target_score = off_target_score
    if off_target_score == 'inf':
      self.off_target_score = 10000

  def serialize_for_display(self):
    """Serialize for the way we are returning json"""
    serialization = {
      "score": self.score,
      "start": self.start,
      "seq": self.seq,
      "PAM": self.PAM,
      "seq_before": self.seq_before,
      "seq_after": self.seq_after,
      "chrom": self.chrom,
      "cut_pos": self.cut_pos,
      "selected": self.selected,
      "has_exome_repeat": self.has_exome_repeat,
      "off_target_score": self.off_target_score,
      "has_functional_domain": self.has_functional_domain
    }
    if self.functional_domain != None:
      serialization["functional_domain"] = self.functional_domain
    return serialization

  def cmp_scheme(self, g):
    return (-g.off_target_score, g.has_functional_domain, g.score)

  def __cmp__(self, other):
    return cmp(self.cmp_scheme(self), self.cmp_scheme(other))

params = {
  "PAM": "NGG",
  "protospacer_len": 20,
  "prime5": True,
  "scoring": "Azimuth",
  "quantity": 100,
  "functional_domains": True,
  "mer_len": 20
}

# azimuth model
print "loading azimuth models", time.time() - start_time
azimuth_saved_model_dir = os.path.join(os.path.dirname(azimuth.__file__), 'saved_models')
model_name = 'V3_model_full.pickle'
azimuth_model_file = os.path.join(azimuth_saved_model_dir, model_name)
with open(azimuth_model_file, 'rb') as f:
  azimuth_model = pickle.load(f)

azimuth_scores_file = 'azimuth_scores_mus.p'
with open(azimuth_scores_file, 'rb') as inp:
  azimuth_scores = pickle.load(inp)

def get_azimuth_score(mer30):
  if mer30 in azimuth_scores:
    return azimuth_scores[mer30]
  else:
    score = azimuth.model_comparison.predict(np.array([mer30]), aa_cut=None, percent_peptide=None, model=azimuth_model, model_file=azimuth_model_file)[0]
    print "generating Azimuth", mer30, score
    azimuth_scores[mer30] = score
    return score

# load in exome
APP_STATIC = "/home/joshm/GUIDES/CRISPR-Library-Designer/static"
exome_seq_path = os.path.join(APP_STATIC, 'data', 'GRCm38_exons')
mer_len = params['mer_len']

# process kmers
# consider all kmers which are followed by NGG
print "preparing mus kmers", time.time() - start_time
exome_mers = {}
for file in os.listdir(exome_seq_path):
  file_loc = os.path.join(exome_seq_path, file)
  with open(file_loc, 'r') as file_data:
    fwdseq = file_data.read()
  revseq = revcom(fwdseq)

  for seq in [fwdseq, revseq]:
    for i in range(len(seq) - mer_len - 2):
      s = seq[i: i + mer_len]
      if seq[i + mer_len + 1 : i + mer_len + 3] != "GG": # only PAMs
        continue
      if 'N' in s:
        continue
      if s in exome_mers:
        exome_mers[s] += 1
      else:
        exome_mers[s] = 1
print 'len(exome_mers) = ', len(exome_mers), time.time() - start_time

# takes in guide OBJECT
# returns whether there is a duplicate in exome
def hasExomeRepeat(protospacer):
  guide_seq = protospacer[-mer_len:] # get PAM-proximal mer_len bases
  hits = exome_mers[guide_seq] # how many times does occur in genome followed by NGG?
  return hits >= 2

# loading CFD preprocessed

#Unpickle mismatch scores and PAM scores
def get_mm_pam_scores():
    try:
        mm_scores = pickle.load(open('mismatch_score.pkl','rb'))
        pam_scores = pickle.load(open('pam_scores.pkl','rb'))
        return (mm_scores,pam_scores)
    except:
        raise Exception("Could not find file with mismatch scores or PAM scores")

#Calculates CFD score
def calc_cfd(wt,sg,pam):
    mm_scores,pam_scores = get_mm_pam_scores()
    score = 1
    sg = sg.replace('T','U')
    wt = wt.replace('T','U')
    s_list = list(sg)
    wt_list = list(wt)
    for i,sl in enumerate(s_list):
        if wt_list[i] == sl:
            score*=1
        else:
            key = 'r'+wt_list[i]+':d'+revcom(sl)+','+str(i+1)
            score*= mm_scores[key]
    score*=pam_scores[pam]
    return (score)

def get_pot_off_targets(seq):
    seq_list = list(seq)
    backup_seq_list = list(seq)
    nts = ['A','T','C','G']
    results = {}
    for a in range(len(seq)):
        for a_sym in nts:
            seq_list[a] = a_sym
            for b in range(a + 1, len(seq)):
                for b_sym in nts:
                    seq_list[b] = b_sym
                    for c in range(b + 1, len(seq)):
                        for c_sym in nts:
                            seq_list[c] = c_sym
                            new_seq = ''.join(seq_list)
                            results[new_seq] = True
                        seq_list[c] = backup_seq_list[c]
                seq_list[b] = backup_seq_list[b]
            seq_list[a] = backup_seq_list[a]
    if seq in results:
        del results[seq]
    return results.keys()

# load preprocessed info
with open("off_target_scores_mus.p", "rb") as inp:
  off_target_scores = pickle.load(inp)

print 'len(off_target_scores) = ', len(off_target_scores), time.time() - start_time

def get_off_target_score(protospacer):
  if hasExomeRepeat(protospacer):
    return 100000
  if not protospacer in off_target_scores:
    score = 0
    off_targets = get_pot_off_targets(protospacer)
    for off_target in off_targets:
        if off_target in exome_mers:
            wt = protospacer + "CGG"
            sg = off_target
            pam = "GG"
            score += exome_mers[off_target] * calc_cfd(wt, sg, pam)
    off_target_scores[protospacer] = score
  return off_target_scores[protospacer]

# Create interval tree for functional domains
print "constructing interval tuples", time.time() - start_time
interval_tuples_dict = {}
ucsc_pfam_f = '../functional_domains/ucsc_pfam_GRCm38.txt'
with open(ucsc_pfam_f, 'r') as pfam_csv:
  csvreader = csv.reader(pfam_csv, delimiter='\t')
  next(csvreader) # skip header

  for row in csvreader:
    chrom = row[1]
    start = row[2]
    end = row[3]