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


  def checkModel(self):

    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

    # sample reweighting
    rewDiagMat = helperFunctions.getReweightDiagMat(self.y, numpy.asmatrix(numpy.unique(numpy.asarray(self.y))))
    rewX = numpy.dot(rewDiagMat,self.X)
    rewY = numpy.dot(rewDiagMat,self.y)

    # calculate w
    self.w = helperFunctions.solveW(rewX, rewY, self.sigmaN)

    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 samples
    self.X = numpy.append(self.X, x, axis=0)
    self.y = numpy.append(self.y, relY, axis=0)

    # sample reweighting
    rewDiagMat = helperFunctions.getReweightDiagMat(self.y, numpy.asmatrix(numpy.unique(numpy.asarray(self.y))))
    rewX = numpy.dot(rewDiagMat,self.X)
    rewY = numpy.dot(rewDiagMat,self.y)

    # calculate w
    self.w = helperFunctions.solveW(rewX, rewY, self.sigmaN, self.w)

    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 reweight(self, alScores, x):

    clsScores = self.infer(x)

    # sample reweighting
    rewDiagMat = helperFunctions.getReweightDiagMat(self.y, numpy.asmatrix(numpy.unique(numpy.asarray(self.y))))
    rewX = numpy.dot(rewDiagMat,self.X)
    var = numpy.sum(numpy.multiply(x,(self.sigmaN*numpy.linalg.solve(numpy.add(numpy.dot(rewX.T,rewX), numpy.identity(x.shape[1])*self.sigmaN), x.T)).T),axis=1)

    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