# -*- coding: utf-8 -*-
import numpy as np
from scipy.stats import entropy
from math import log10

def sense(world, p, z, pHit, pMiss):
	q = np.array(p)
	world = np.array(world)

	q[world==z] *= pHit
	q[world!=z] *= pMiss

	q = normalize(q)

	return q

def normalize(p):
	q = np.array(p)
	q /= np.sum(q)
	return q

def move(p, U, pUndershoot, pExact, pOvershoot):
	n = []

	U = U % len(p)
	# Shift array
	q = np.append(p[-U:], p[:-U])

	for i, item in enumerate(q):
		s = item * pExact
		s += q[i-1] * pOvershoot
		s += q[(i+1)%len(q)] * pUndershoot
		n.append(s)
	return np.array(n)

def entropys(p):
    s = 0
    for x in p:
        s += x*log10(x)
    return -s

p = np.array([0.2] * 5)
world = ['green', 'red', 'red', 'green', 'green']
measurements = ['red', 'red']
motions = [1, 1]

# Sense parameters
pHit = 0.6
pMiss = 0.2
# Inexact motion parameters
pExact = 0.8
pOvershoot = 0.1
pUndershoot = 0.1

# Sense multiple data
for z in measurements:
	p = sense(world, p, z, pHit, pMiss)

# Move

p = [0, 1, 0, 0, 0]
p = move(p,1, pUndershoot, pExact, pOvershoot)
p = move(p,1, pUndershoot, pExact, pOvershoot)
print(p)

# Limit distribution
for _ in range(1000):
	p = move(p,1, pUndershoot, pExact, pOvershoot)
print(p)

# Sense and move
p = np.array([0.2] * 5)

print("Entropy: {}".format(entropys(p)))

p = [0, 1., 0, 0, 0]

for u, z in zip(motions, measurements):
	p = sense(world, p, z, pHit, pMiss)
	p = move(p, u, pUndershoot, pExact, pOvershoot)

print(p)
print("Entropy: {}".format(entropys(p)))