#! /usr/bin/python

import numpy
import scipy.special

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.00178,
               configFile=None):

    self.sigmaN = helperFunctions.getConfig(configFile, 'reweighting', 'sigmaN', sigmaN, 'float', True)

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


  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 = X*X.T
    self.alpha = numpy.linalg.solve(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.X*x.T
    selfK = numpy.sum(numpy.multiply(x,x), axis=1)

    # get score of new sample
    infY = self.infer(x, k=k)

    # 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.float32) * -1.0
    term2[0:self.X.shape[0],:] = numpy.linalg.solve(self.K + numpy.identity(self.X.shape[0], dtype=numpy.float32)*self.sigmaN, k)

    term3 = infY - relY

    self.alpha = numpy.append(self.alpha, numpy.zeros((1,self.alpha.shape[1])), axis=0) + 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, relY, axis=0)

    self.checkModel()


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

    if k is None:
      k = self.X*x.T

    pred = numpy.asmatrix(k.T*self.alpha)

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

    return pred


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

    if k is None:
        k = self.X*x.T

    if selfK is None:
        selfK = numpy.sum(numpy.multiply(x,x), axis=1)

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

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

    return sigmaF


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

    if k is None:
      k = self.X*x.T

    clsScores = self.infer(x,k)
    var = self.calcSigmaF(x,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