#!/usr/bin/env python
# encoding: utf-8
"""
Predictor.py

Created by Wim Vranken on 25th July 2018, based on original DynaMine code (Elisa Cilia)
Copyright (c) 2018 Wim Vranken. All rights reserved.
"""

import os
import glob
import numpy

from b2bTools.general import Util as b2bUtil
from b2bTools.general.Io import B2bIo


# TODO: include file/JSON writing class in here for default file output options
class DynaMine(B2bIo):

  version = '4.0'
  # scriptName = "b2bTools.singleSeq.DynaMine.Predictor"
  scriptName = "python.b2bTools.singleSeq.DynaMine.Predictor"
  name = "DynaMine suite"

  aas = '-ACDEFGHIKLMNPQRSTVWYX'

  model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), "models/")

  def __init__(self,predictionTypes=('backbone', 'sidechain', 'ppII', 'coil', 'sheet', 'helix')):

    """
    :param predictionTypes: a list of the prediction types that need to be performed. By default all possible types are done.
    """

    self.windowsize = 51

    self.modelDict = {}
    self.modelBias = {}
    self.input = False
    self.outputfilenames = []
    self.allPredictions = {}

    self.types_ = []
    for predictionType in predictionTypes:
      if predictionType in ('backbone','sidechain'):
        suffix = ''
      else:
        suffix = '_new'
      self.types_.append((predictionType,suffix))

    self.loadModels()
    self.adjustModels()

    # Additional info for writing files
    self.references = ['doi: 10.1038/ncomms3741 (2013)', 'doi: 10.1093/nar/gku270 (2014)']
    self.infoTexts = ['Generated by the DynaMine suite','See http://bio2byte.be']

    self.informationPerPredictor = {
      self.name: {
          'references': self.references,
          'info':    ";".join(self.infoTexts),
          'version': self.version,
          'origin':  self.scriptName
                }
          }

  def loadModels(self):
    modelLoaded = True
    for (type_,suffix) in self.types_:

      pickleWindowList = []
      modelDir = os.path.join(self.model_path, type_)
      pickleFileNames = glob.glob(os.path.join(modelDir,"read*.model{}.pkl.bz2".format(suffix)))
      for pfn in pickleFileNames:
        pickleWindowList.append((int(os.path.split(pfn)[1].split('.')[0][4:]),pfn))

      ### Doesn't find pklfile
      for (w,pklfile) in pickleWindowList:
        #print w, pklfile
        if os.path.isfile(pklfile):
          p = b2bUtil.Pickle()
          model = p.load_data(pklfile)
          modelKey = (type_, w)
          self.modelDict[modelKey] = []
          self.modelBias[modelKey] = model['bias']
          # Precalculate for faster retrieval
          for predPos in list(model.keys()):
            if not predPos == 'bias':
              self.modelDict[modelKey].append({})
              for aa in self.aas:
                self.modelDict[modelKey][-1][aa] = 0.0
              for (correction, aaCodes) in model[predPos]:
                for aaCode in aaCodes:
                  self.modelDict[modelKey][-1][aaCode] += correction
      else:
        modelLoaded = False
    return modelLoaded

  def adjustModels(self):
    """ Substitute the values for '-' with the median over the position
    """
    for modelKey, model in self.modelDict.items():
      # w = len(model) / 2
      for i in range(len(model)):
        newdistribution = self.filterOutliers(list(model[i].values())) # remove outliers
        #newdistribution = model[i].values()
        m = numpy.median(newdistribution) # compute the median per position
        # print "before:", model[i]['-'], "after:", m
        model[i]['-'] = m

  # Warning: THIS IS NOW NOT WORKING FROM A DICT BUT FROM A LIST OF TUPLES!
  def predictSeqs(self, seqs, force_windowsize=None):

    """
    :param seqs: List/tuple of (seqId,sequenceString) tuples
    :param force_windowsize: Force a particular sequence window size when running
    :return: True if all sequences predicted, False if problems

    See self.allPredictions for the predictions themselves.
    """

    count_short = 0

    for (seqId,sequence) in seqs:

      print('Predicting {}...\n'.format(seqId), end=' ')

      seqLen = len(sequence)
      self.allPredictions[seqId] = {}

      if force_windowsize:
        self.windowsize = force_windowsize
      else:
        if seqLen > 51:
          self.windowsize = 51
        elif seqLen >= 25:
          self.windowsize = 25
        elif seqLen >= 19:
          self.windowsize = 19
        elif seqLen == 5:
          self.windowsize = 5
        elif seqLen < 5:
          self.windowsize = 5
          count_short += 1
          print('Skipping, too short.')
          continue
        else:
          if seqLen % 2 == 0:
            self.windowsize = seqLen - 1
          else:
            self.windowsize = seqLen

      for (predictionType,suffix) in self.types_:
        self.allPredictions[seqId][predictionType] = self.computePredictions(sequence, predictionType)

    if count_short == len(seqs):
      return False
    else:
      return True

  def computePredictions(self, s, type_):
    """
    Returns a list of lists (one for each amino acid in the sequence) containing two strings ['aa type one-letter code', 'prediction']
    """
    # Wim speedup attempt
    emptyTailLen = int((self.windowsize - 1) / 2)
    s_app = self.correctFragment("-" * emptyTailLen + "".join([x[0] for x in s]) + "-" * emptyTailLen)
    #print s_app
    modelKey = (type_, self.windowsize)
    seqLen = len(s_app)
    preds = []

    for seqPos in range(seqLen):
      preds.append(self.modelBias[modelKey])

    for seqPos in range(seqLen):
      aaCode = s_app[seqPos]
      predPosBase = emptyTailLen + seqPos
      for otherSeqPos in range(max(0,seqPos - emptyTailLen),min(seqLen,seqPos+emptyTailLen+1)):
        predPos = predPosBase - otherSeqPos
        preds[otherSeqPos] += self.modelDict[modelKey][predPos][aaCode]

    for seqPos in range(seqLen):
      preds[seqPos] = (s_app[seqPos],preds[seqPos])

    return preds[emptyTailLen:-emptyTailLen]

  #
  # Functions ripped from other scripts in original dynaMine distr
  #

  def correctFragment(self,frag):
    seq = ''
    for j in frag:
      if not j[0] in self.aas:
        seq += 'X'
      else:
        seq += j[0]
    return seq

  def filterOutliers(self,values):
    q3 = numpy.percentile(values, 75, interpolation='higher')
    q1 = numpy.percentile(values, 25, interpolation='lower')
    iqr = q3 - q1
    newdistribution = [v for v in values if v >= q1 - 1.5 * iqr and v <= q3 + 1.5 * iqr]
    return newdistribution
