#! /usr/bin/python

import numpy
import scipy.special
import sklearn.kernel_ridge

import sys
import os
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)),os.pardir))
import helperFunctions

class Reweighter:

    def __init__(self, sigmaN=0.01, gamma=None, kernel='rbf', configFile=None):

        self.sigmaN = helperFunctions.getConfig(configFile, 'activeLearning', 'sigmaN', sigmaN, 'float', True)
        self.kernel = helperFunctions.getConfig(configFile, 'activeLearning', 'kernel', kernel, 'str', True)
        self.gamma = helperFunctions.getConfig(configFile, 'activeLearning', 'gamma', gamma, 'float', True)
        self.numKernelCores = helperFunctions.getConfig(configFile, 'activeLearning', 'numKernelCores', 1, 'int', True)

        self.K = []
        self.alpha = []
        self.X = []
        self.y = []

        self.allowedKernels = ['rbf', 'sigmoid', 'polynomial', 'poly', 'linear', 'cosine']

        if self.kernel not in self.allowedKernels:
            raise Exception('Unknown kernel %s!'%self.kernel)


    def kernelFunc(self, x1, x2=None, gamma=None):

        if gamma is not None:
            self.gamma = gamma

        if self.kernel in ['rbf']:
            return numpy.asmatrix(sklearn.kernel_ridge.pairwise_kernels(x1, x2, metric=self.kernel, gamma=self.gamma, n_jobs=self.numKernelCores), dtype=numpy.float)
        else:
            return numpy.asmatrix(sklearn.kernel_ridge.pairwise_kernels(x1, x2, metric=self.kernel, n_jobs=self.numKernelCores), dtype=numpy.float)


    def checkModel(self):

        if not numpy.all(numpy.isfinite(self.K)):
            raise Exception('not numpy.all(numpy.isfinite(self.K))')

        if not numpy.all(numpy.isfinite(self.alpha)):
            raise Exception('not numpy.all(numpy.isfinite(self.alpha))')

        if not numpy.all(numpy.isfinite(self.X)):
            raise Exception('not numpy.all(numpy.isfinite(self.X))')

        if not numpy.all(numpy.isfinite(self.y)):
            raise Exception('not numpy.all(numpy.isfinite(self.y))')


    # X.shape = (number of samples, feat dim), y.shape = (number of samples, 1)
    def train(self, X, y):

        self.X = X
        self.y = numpy.asmatrix(y == -1, dtype=numpy.int)*2 - 1

        self.K = self.kernelFunc(X,X)
        self.alpha = numpy.linalg.solve(numpy.add(self.K, numpy.identity(self.X.shape[0], dtype=numpy.float32))*self.sigmaN, self.y)

        self.checkModel()


    # x.shape = (1, feat dim), y.shape = (1, 1)
    def update(self, x, y):

        # relabel
        relY = numpy.asmatrix(y == -1, dtype=numpy.int)*2 - 1

        # get new kernel columns
        k =  self.kernelFunc(self.X, x)
        selfK =  self.getSelfK(x)

        # update alpha
        term1 = 1.0 / (self.sigmaN + self.calcSigmaF(x, k, selfK).item(0))

        term2 = numpy.asmatrix(numpy.ones((self.X.shape[0] + 1,x.shape[0])), dtype=numpy.float) * (-1.0)
        term2[0:self.X.shape[0],:] = numpy.linalg.solve(numpy.add(self.K, numpy.identity(self.X.shape[0], dtype=numpy.float)*self.sigmaN), k)

        term3 = self.infer(x, k=k) - relY

        self.alpha = numpy.add(numpy.append(self.alpha, numpy.zeros((1,self.alpha.shape[1])), axis=0), numpy.dot(numpy.dot(term1,term2),term3))

        # update K
        self.K = numpy.append(numpy.append(self.K, k, axis=1), numpy.append(k.T, selfK, axis=1), axis=0)

        # update samples
        self.X = numpy.append(self.X, x, axis=0)
        self.y = numpy.append(self.y, y, axis=0)

        self.checkModel()


    # X.shape = (number of samples, feat dim), loNoise = {0,1}
    def infer(self, x, loNoise=False, k=None):

        if k is None:
            k = self.kernelFunc(self.X, x)

        loNoise = loNoise and (self.yUni == -1).any()

        pred = numpy.asmatrix(numpy.dot(k.T,self.alpha[:,int(loNoise):]))

        if not numpy.all(numpy.isfinite(pred)):
            raise Exception('not numpy.all(numpy.isfinite(pred))')

        return pred


    def calcSigmaF(self, x, k=None, selfK=None):

        if k is None:
            k = self.kernelFunc(self.X, x)

        if selfK is None:
            selfK = self.getSelfK(x)

        sigmaF = numpy.subtract(selfK, numpy.sum(numpy.multiply(numpy.linalg.solve(numpy.add(self.K, numpy.identity(self.X.shape[0], dtype=numpy.float)*self.sigmaN), k),k).T, axis=1))

        if not numpy.all(numpy.isfinite(sigmaF)):
            raise Exception('not numpy.all(numpy.isfinite(sigmaF))')

        return sigmaF


    def getSelfK(self, x):

        selfK = numpy.asmatrix(numpy.empty([x.shape[0],1], dtype=numpy.float))
        for idx in range(x.shape[0]):
            selfK[idx,:] = self.kernelFunc(x[idx,:])

        return selfK


    # X.shape = (number of samples, feat dim)
    def reweight(self, alScores, x, k=None):

        if k is None:
            k = self.kernelFunc(self.X, x)

        clsScores = self.infer(x, k=k)
        var = self.calcSigmaF(x, k=k)

        probs = 0.5 - 0.5*scipy.special.erf(clsScores*(-1)/numpy.sqrt(2*var))
        newAlScores = numpy.multiply(alScores,(1 - probs))

        if not numpy.all(numpy.isfinite(newAlScores)):
            raise Exception('not numpy.all(numpy.isfinite(newAlScores))')

        return newAlScores