# pip3 install --user http://bit.ly/csc161graphics
from graphics import *
import time
import numpy as np
import math
from map import *
import random
import colorsys

class VirtualRobot:
    def __init__(self, worldmap):
        self.x = np.random.rand()*worldmap.width
        self.y = np.random.rand()*worldmap.height
        self.heading = math.pi*2*np.random.rand()
        self.worldmap = worldmap

    def move(self, distance):
        self.x = self.x + (math.cos(self.heading) * distance)
        self.y = self.y + (math.sin(self.heading) * distance)

    def rotate(self, radians):
        self.heading += float(radians)
        self.heading %= 2 * math.pi

    def read(self):
        return self.worldmap.white(self.x, self.y)

# ------------------------------------------------------------------------
class Particle:
    def __init__(self, worldmap):
        self.weight = 1
        self.x = np.random.rand()*worldmap.width
        self.y = np.random.rand()*worldmap.height
        self.heading = math.pi*2*np.random.rand()
        self.worldmap = worldmap

    def move(self, rot, distance):
        if np.random.rand() < 0.98:
            self.translate(distance)
            self.rotate(rot)
        else:
            self.x = np.random.rand()*self.worldmap.width
            self.y = np.random.rand()*self.worldmap.height
            self.heading = math.pi*2*np.random.rand()

    def translate(self, distance):
            distance = float(distance) + random.gauss(0.0, 0.4*distance)
            self.x = self.x + (math.cos(self.heading) * distance)
            self.y = self.y + (math.sin(self.heading) * distance)

    def rotate(self, radians):
        self.heading += float(radians) + random.gauss(0.0, 0.1*radians+0.02)
        self.heading %= 2 * math.pi

    def read(self):
        return self.worldmap.white(self.x, self.y)


    def __repr__(self):
        return "(%f, %f)" % (self.x, self.y)

def compute_mean_point(particles):
    m_x, m_y, m_count = 0, 0, 0
    for p in particles:
        m_count += p.weight
        m_x += p.x * p.weight
        m_y += p.y * p.weight

    if m_count == 0:
        return -1, -1

    m_x /= m_count
    m_y /= m_count

    return m_x, m_y

def low_variance_sampler (particles, worldmap):
    weights = []
    # normalize
    for p in particles:
        weights.append(p.weight)
    weights = weights / np.sum(weights) 
    # print(np.sum(weights))

    new_particles = []

    M = len(particles)
    r = np.random.uniform(0, 1.0/M)
    c = weights[0]
    i = 0

    for m in range (0, M):
        # try:
        u = r + m * 1.0/M
        while u > c:
            i = i + 1
            c = c + weights[i]
        p = Particle(worldmap)
        p.x = particles[i].x
        p.y = particles[i].y
        p.heading = particles[i].heading
        new_particles.append(p)

        # except IndexError:
        #     new_particles.append(Particle(worldmap))
    return new_particles

def particle_filter(particles, read, move, rot, worldmap):
    new_particles = []
    M = len(particles)

    for p in particles:
        p.move(rot, move)
        # w = p(z|x) probabilidade de ter lido uma cor estadando em x
        if worldmap.inside(p.x, p.y):
            if p.read() == read:
                p.weight = 0.80
            else:
                p.weight = 0.20
            # p_d = p.read_sensor(world)
            # p.w = w_gauss(r_d, p_d)
        else:
            ptmp = Particle(worldmap)
            p.x = ptmp.x
            p.y = ptmp.y
            p.heading = ptmp.heading
            p.weight = 1.e-300
    # print(particles)
    return low_variance_sampler(particles, worldmap)
    


# ------------------------------------------------------------------------
def colorFix(floatcolor):
    r = int(floatcolor*240)
    if r < 0:
        r = 0
    elif r > 255:
        r = 255
    return r


def main():
    width = 3370.39
    height = 2383.94
    zoom = 2.5
    worldmap = Map(width, height, "../3.txt")
    robot = VirtualRobot(worldmap)

    PARTICLE_COUNT = 2000
    particles = []

    win = GraphWin('Particle', width/zoom, height/zoom, autoflush=False)
    win.setBackground('black')

    # create particles
    for i in range(0, PARTICLE_COUNT):
        particles.append(Particle(worldmap))

    while True:
        worldmap.draw(win, zoom)
        rot = 0
        move = 50

        # # ---------- Try to find current best estimate for display ----------
        m_x, m_y = compute_mean_point(particles)
        # movimento
        if robot.x < width*0.1 or width*0.9 < robot.x:
            rot = np.random.uniform(math.pi/2,math.pi/2)
            robot.rotate(rot)
        if robot.y < height*0.1 or height*0.9 < robot.y:
            rot = np.random.uniform(math.pi/2,math.pi/2)
            robot.rotate(rot)
        robot.move(move)

        #leitura
        read = robot.read()

        particles = particle_filter(particles, read, move, rot, worldmap)

        # ---------- display particles ----------
        for p in particles:
            pt = Point(p.x/zoom, p.y/zoom)
            h = p.heading/(math.pi*2)
            rgb = colorsys.hsv_to_rgb(h, 0.8, 0.8)
            rgb = color_rgb(colorFix(rgb[0]), colorFix(rgb[1]), colorFix(rgb[2]))
            cir = Circle(pt, 10/zoom)
            cir.setWidth(0)
            cir.setFill(rgb)
            cir.draw(win)


        # ---------- show mean point ----------
        pt = Point(m_x/zoom, m_y/zoom)
        cir = Circle(pt, 30/zoom)
        cir.setFill('red')
        cir.draw(win)

        # ---------- display robot ----------
        h = robot.heading/(math.pi*2)
        rgb = colorsys.hsv_to_rgb(h, 0.8, 0.8)
        rgb = color_rgb(colorFix(rgb[0]), colorFix(rgb[1]), colorFix(rgb[2]))
        pt = Point(robot.x/zoom, robot.y/zoom)
        cir = Circle(pt, 30/zoom)
        cir.setFill(rgb)
        cir.draw(win)

        win.update()
        win.delete('all')

main()