from __future__ import print_function

__author__ = 'Frank Sehnke, sehnke@in.tum.de'


from scipy import random, zeros, ones, exp, sqrt, cos, log

stND = zeros(1000)
for i in range(1000):
    x = -4.0 + float(i) * 8.0 / 1000.0
    stND[i] = 1.0 / 2.51 * exp(-0.5 * (x) ** 2)


# FIXME: different class name?
class MixtureOfGaussians:
    def __init__(self, numOGaus=10, alphaA=0.02, alphaM=0.02, alphaS=0.02):
        self.alphaA = alphaA
        self.alphaM = alphaM
        self.alphaS = alphaS
        self.minSig = 0.000001
        self.numOGaus = numOGaus #Number of Gaussians
        self.rangeMin = -20.0
        self.rangeMax = 20.0
        self.epsilon = (self.rangeMax - self.rangeMin) / (sqrt(2.0) * float(self.numOGaus - 1)) #Initial value of sigmas

        self.propFakt = 1.0 / float(self.numOGaus)
        self.distFakt = 1.0 / float(self.numOGaus - 1)
        self.distRange = self.rangeMax - self.rangeMin

        self.sigma = ones(self.numOGaus)
        self.mue = zeros(self.numOGaus)
        self.alpha = ones(self.numOGaus)
        self.sigma *= self.epsilon
        self.alpha /= float(self.numOGaus)
        for i in range(self.numOGaus):
            self.mue[i] = self.distRange * float(i) * self.distFakt + self.rangeMin
        self.baseline = 0.0
        self.best = 0.000001

    def getStND(self, x, mue=0.0, sig=1.0):
        x = (x - mue) / sig
        if abs(x) >= 4.0: return 0.000000001
        x = int((x + 4.0) / 8.0 * 1000)
        return stND[x] / sig

    #get mixture for point x
    def getGaus(self, alpha, mue, sigma, x):
        dens = zeros(self.numOGaus)
        for g in range(self.numOGaus):
            dens[g] = self.getStND(x, mue[g], sigma[g])
        return sum(self.sigmo(alpha) * dens)

    def drawSample(self):
        sum = 0.0
        rndFakt = random.random()
        for g in range(self.numOGaus):
            sum += self.sigmo(self.alpha[g])
            if rndFakt < sum:
                if self.sigma[g] < self.minSig: self.sigma[g] = self.minSig
                x = random.normal(self.mue[g], self.sigma[g])
                break
        return x

    def learn(self, x, y):
        #learning overalls
        norm = zeros(self.numOGaus)
        nOver = self.getGaus(self.alpha, self.mue, self.sigma, x)
        for g in range(self.numOGaus):
            norm[g] = self.getStND(x, self.mue[g], self.sigma[g]) / nOver
        self.baseline = 0.99 * self.baseline + 0.01 * y
        if y > self.best: self.best = y
        fakt = (y - self.baseline) / (self.best - self.baseline)
        if fakt < -1.0: fakt = -1.0

        #alpha learning
        sigmoA = self.sigmo(self.alpha)
        self.alpha += self.alphaA * self.propFakt * sigmoA * (1.0 - sigmoA) * fakt * norm #(1.0-sigmoA)*
        sigmoA = self.sigmo(self.alpha)
        sigmoA /= sum(sigmoA)
        self.alpha = self.invSigmo(sigmoA)

        #mue learning
        sigmoA = self.sigmo(self.alpha)
        self.mue += self.alphaM * fakt * (x - self.mue) * sigmoA * norm

        #sigma learning
        if fakt > 0.0:
            self.sigma += self.alphaS * fakt * ((x - self.mue) ** 2 - self.sigma ** 2) / self.sigma * sigmoA * norm

    def sigmo(self, a):
        return 1.0 / (1.0 + exp(-1.0 * a))

    def invSigmo(self, a):
        return - log(1.0 / a - 1.0)

    def getSample(self):
        sampleX = self.drawSample()
        return sampleX

if __name__ == '__main__':
    m = MixtureOfGaussians()
    for i in range(10000):
        x = m.getSample()
        n = x / m.distRange * 10.0
        if abs(n) > 5.0: n = 0.0
        y = (20.0 + n ** 2 - 10.0 * cos(2.0 * 3.1416 * n)) / 55.0 + random.normal(0, 0.2) #one dimensional rastrigin
        m.learn(x, y)
    print((m.alpha))
    print((m.mue))
    print((m.sigma))

