# class ParticleFilterPoseEstimator
# O. Bittel; 31.10.2014

from math import *
import numpy.matlib as ml
import numpy as np
import random as rd
import matplotlib.pyplot as plt

class ParticleFilterPoseEstimator:

    # --------
    # initializes particle filter.
    #
    def __init__(self):
        self.nPart = 100 # number of particles
        self.chi = [[0.0,0.0,0.0] for i in range(self.nPart)] # particle set
        self.mean = [0.0,0.0,0.0]
        self.covariance = ml.zeros((3,3))
        self.T = 0.01 # time step
        self.room = None # room for particles
        self.measuringDiversity = False # If true the number of different particles are counted

    # --------
    # print robot attributes.
    #
    def __repr__(self):
        return 'Particle filter pose: [x=%6.4f y=%6.4f orient=%6.4f]' % (self.x, self.y, self.orientation)

    # --------
    # set time step.
    #
    def setTimestep(self, T):
        self.T = float(T)

    def setRoom(self, testroom):
        self.room = testroom


    # --------
    # initialize the particle filter with n particles following
    # an uniform distribution in the interval [poseFrom, poseTo].
    #
    def initialize(self, poseFrom, poseTo, n = 100):
        self.chi = [[0.0,0.0,0.0] for i in range(n)]
        for i in range(n):
            self.chi[i][0] = poseFrom[0] + (poseTo[0]-poseFrom[0])*rd.random()
            self.chi[i][1] = poseFrom[1] + (poseTo[1]-poseFrom[1])*rd.random()
            self.chi[i][2] = poseFrom[2] + (poseTo[2]-poseFrom[2])*rd.random()
        self.nPart = n
        self.__updateMeanCov()
        return


    def __updateMeanCov(self):
        #print '__updateMeanCov: ', self.mean
        x = [0.0 for i in range(len(self.chi))]
        y = [0.0 for i in range(len(self.chi))]
        theta = [0.0 for i in range(len(self.chi))]
        i = 0
        for p in self.chi:
            x[i] = p[0]
            y[i] = p[1]
            theta[i] = p[2]
            i += 1
        mean_theta = self.__mean_orientation(theta)
        self.mean = [np.mean(x),np.mean(y),mean_theta]
        self.covariance = ml.zeros((3,3))
        self.covariance[0:2,0:2] =  np.cov([x,y])
        self.covariance[2,2] = self.__variance_orientation(mean_theta,theta)
        #print '__updateMeanCov: ', self.mean
        #print "x:     ", self.mean[0], self.covariance[0,0], x
        #print "y:     ", self.mean[1], self.covariance[1,1], y
        #print 'theta_mean: %20.18f theta_sigma: %20.18f' % (self.mean[2]*180/pi, sqrt(self.covariance[2,2])*180/pi)
        #print "theta: ", theta
        return


    # returns mean of orientations.
    # mean can be ambigious (e.g. theta = [0, pi/2, pi, 3*pi/2].
    # We use a heuristic: the mean never lies between two consecutive orientation with the largest angular difference.
    # In some few cases the heuristic leads to a wrong mean.
    def __mean_orientation(self, theta):
        n = len(theta)
        if n == 1:
            return theta[0]
        theta.sort()
        #print theta
        d = theta[0] - theta[n-1] + 2*pi
        min = n-1
        max = 0
        for i in range(0,n-1):
            if theta[i+1] - theta[i] > d:
                d = theta[i+1] - theta[i]
                min = i
                max = i+1
        # print min, max
        # print theta[min], theta[max]
        # Mittelwert liegt zwischen theta[max] und theta[min]:
        sum = 0.0
        for i in range(n):
            sum += (theta[i] - theta[max])
        mean = (sum/n + theta[max]) % (2*pi)
        return mean


    def __variance_orientation(self, mean, theta):
        n = len(theta)
        if n == 1:
            return 0.0
        sum = 0.0
        for i in range(n):
            diff = (theta[i] - mean + pi) % (2*pi) - pi # orientation difference from [-pi, pi)
            sum += diff**2
        return sum/n


    # --------
    # get the current robot position estimate.
    #
    def getPose(self):
        return self.mean


    # --------
    # get the current position covariance.
    #
    def getCovariance(self):
        return self.covariance


    # --------
    # get particles.
    #
    def getParticles(self):
        return self.chi


    # --------
    # integrate motion = [v, omega] and covariance sigma_motion for the next time step
    #
    def integrateMovement(self, motion, sigma_motion):

        for i in range(len(self.chi)):
            # Sample noisy motion:
            v = rd.gauss(motion[0],sqrt(sigma_motion[0,0]))
            omega = rd.gauss(motion[1],sqrt(sigma_motion[1,1]))
            # Apply noisy motion to particle chi[i]:
            theta = self.chi[i][2]
            self.chi[i][0] += v*self.T*cos(theta) # x
            self.chi[i][1] += v*self.T*sin(theta) # y
            self.chi[i][2] = (self.chi[i][2] + omega*self.T) % (2*pi) # orientation

        self.__updateMeanCov()
        return


    # --------
    # integrate motion = [v, omega] and covariance sigma_movement for the next time step
    #
    def integrateMeasurement(self, z, sigma_z):
        # z = number of measurements || should be same as the number of landmarks'

        # Weighting:
        nPart = len(self.chi)
        w = [0.0] * nPart
        wsum = [0.0] * nPart
        for i in range(nPart):
            w[i] = self.getWeight(self.chi[i], z, sigma_z)
            if i == 0:
                wsum[i] = w[i]
            else:
                wsum[i] = wsum[i-1] + w[i]

        # Resampling:
        chi_old = self.chi
        self.chi = [[0.0,0.0,0.0] for i in range(nPart)]

        if self.measuringDiversity:
            hist = [0 for i in range(nPart)]
            ndiff = 0;

        for i in range(nPart):
            r = rd.random()*wsum[nPart-1]
            # roulette method with binary search:
            li = 0;
            re = nPart-1;
            while li <= re:
                m = int((li + re)/2)
                if r < wsum[m]:
                    re = m-1
                elif r > wsum[m]:
                    li = m+1
                else:
                    break
            # Vorsicht!!!!
            # Referenzzuweisung waere falsch:
            # self.chi[i] = chi_old[m]
            self.chi[i][0] = chi_old[m][0]
            self.chi[i][1] = chi_old[m][1]
            self.chi[i][2] = chi_old[m][2]

            if self.measuringDiversity:
                hist[m] += 1
                if hist[m] == 1:
                    ndiff += 1

        if self.measuringDiversity:
            print('Anz. unterschiedl. Partikel: ', ndiff, ' Gewichtssumme: ', wsum[nPart-1])

        self.__updateMeanCov()
        return


    def getWeight(self,particle, z, sigma_z):
        prob = self.room.sense((particle[0], particle[1]), particle[2], False)
        return prob


    # --------
    # 1-dimensional normal distribution N(mu,sigma2)
    #
    def __ndf(self,x,mu,sigma2):
        c = 1/(sqrt(2*sigma2*pi))
        return c*np.exp(-0.5 * (x-mu)**2/sigma2)