#!/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 time
import numpy
from packaging import version

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:
            tic = time.perf_counter()

            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
                )

            toc = time.perf_counter()
            self.allPredictions[seqId]["dynamine_execution_time"] = toc - tic

        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):
        if version.parse(numpy.__version__) < version.parse('1.22'):
            q3 = numpy.percentile(values, 75, interpolation="higher")
            q1 = numpy.percentile(values, 25, interpolation="lower")
        else:
            q3 = numpy.percentile(values, 75, method="higher")
            q1 = numpy.percentile(values, 25, method="lower")

        iqr = q3 - q1
        newdistribution = [
            v for v in values if v >= q1 - 1.5 * iqr and v <= q3 + 1.5 * iqr
        ]

        return newdistribution
