#! /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.invCreg = []
    self.w = []  # .shape = (feat dim, number of unique classes)
    self.X = []
    self.y = []


  def checkModel(self):

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

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

    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):

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

    # calculate inverse of C_reg
    self.invCreg = numpy.linalg.inv((self.X.T*self.X) + numpy.identity(self.X.shape[1])*self.sigmaN)

    # calculate w
    self.w = self.invCreg*self.X.T*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

    # update w and C_reg
    tmpVec = self.invCreg*x.T
    tmpScalar = 1.0 + x*tmpVec
    self.w = self.w + tmpVec*((relY - x*self.w)/tmpScalar)
    self.invCreg = self.invCreg - ((tmpVec*tmpVec.T)/tmpScalar)

    # 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):

    pred = numpy.asmatrix(x*self.w)

    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):

    sigmaF = numpy.sum(numpy.multiply(x,(self.sigmaN*self.invCreg*x.T).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):

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

    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