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

Created by Wim Vranken on 1st May 2020, as wrapper around original PSPer code
Note that the original DisoMine code was upgraded to Python3.7 and simplified by taking out redundant code for production version
The DisoMine code was also modified to better work in conjunction with the data structures used here.
Copyright (c) 2020- Wim Vranken. All rights reserved.
"""

import time
from b2bTools.general.Io import B2bIo
from b2bTools.singleSeq.PSPer.standalone import standalone
import numpy as np

class PSPer(B2bIo):

  version = '1.0'
  scriptName = "b2bTools.singleSeq.PSPer.Predictor"
  name ="PSPer"

  def __init__(self):
    # Additional info for writing files
    self.references = ['Not yet available, submitted']
    self.infoTexts = ['Generated by PSPer', 'See http://bio2byte.be']
    self.allPredictions = {}
    self.informationPerPredictor = {
      self.name: {
        'references': self.references,
        'info':       ";".join(self.infoTexts),
        'version':    self.version,
        'origin':     self.scriptName
      }
    }

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

      if not psper_results:
          print("PSPer failed....")
          return
      elif 'error' in psper_results:
          print("PSPer failed with known error....", psper_results['error'])
          return

      psp_results = {}
      sequences_count = len(psper_results['results'])

      for seq_index in range(sequences_count):
            tic = time.perf_counter()

            seq_dict = psper_results['results'][seq_index]
            key = seq_dict['proteinID']
            del seq_dict['proteinID']

            sequence = [residue for residue in seq_dict['sequence']]
            del seq_dict['sequence']

            psp_results[key] = { 'seq': sequence }

            for prediction_key, prediction_value in seq_dict.items():
                if type(prediction_value) == list or type(prediction_value) == np.array or type(prediction_value) == np.ndarray:
                    psp_results[key][prediction_key] = [(residue, prediction) for (residue, prediction) in zip(sequence, prediction_value)]
                else:
                    psp_results[key][prediction_key] = prediction_value

            toc = time.perf_counter()
            psp_results[key]['psper_execution_time'] = psp_results[key]['psper_execution_time'] + (toc - tic)

      self.allPredictions = psp_results
