"""
Very simple particle filter demo corresponding to CS354 Homework
Exercises.  Tracks a robot moving in a toroidal grid world.

Author: Nathan Sprague &&
"""
import random

class Particle:
    """A single particle containing a weight and a position estimate."""

    # class constants:
    NUM_ROOMS = 4
    MOVE_PROB = .5 # Moves or stays still.
    SENSOR_ACCURACY = .8 # can be off by one to the left or right.

    def __init__(self, w, x):
        """ Particle constructor.
        Arguments:
        w - weight for this particle (float)
        x - state for this particle (int)
        """

        self.weight = w
        self.x = x

    def move(self, u):
        """Update the state of this particle probabilisticaly according to the
           motion mode.

        Arguments:
        u - The action. either +1 (move right) or -1 (move left)

        """
        if random.random() > self.MOVE_PROB:
            self.x = (self.x + u) % self.NUM_ROOMS

    def sense(self, z):
        """Update the weight assigned to this particle according to the
        sensor model. If the sensor reading agrees with the particle
        state, the weight will tend to increase.

        Arguments:
            z - The output of a "room detector". It will be between 0 and
                NUM_ROOMS -1.
        """

        if z == self.x:
            self.weight = self.SENSOR_ACCURACY * self.weight
        elif (z == ((self.x + 1) % self.NUM_ROOMS) or
              z == ((self.x - 1) % self.NUM_ROOMS)):
            self.weight = ((1 - self.SENSOR_ACCURACY) / 2.0) * self.weight
        else:
            self.weight = 0.0


    def copy(self):
        """Return a copy of this particle."""
        return Particle(self.weight, self.x)

def normalize_particles(particle_list):
    """Update weights so that they sum to 1.

    Argument: particle_list - A python list containing Particle
    objects

    No return value.

    """
    #UNFINISHED
    pass


def calc_probability(particle_list, x):
    """Return the estimated probability that the robot is in state x.

    Arguments:
        particle_list - A (normalized) list of particles.
        x             - The room to check (int)

    Returns:
        The probability that the robot is in state x (float)
    """

    #UNFINISHED
    return 0

def resample(particle_list):
    """Resample the particles.  Particles in the resulting list will all
    have the same weight.

    Algorithm from Table 4.4 in Thrun et. al. Probalistic Robotics.

    Argument: 
    particle_list - A python list containing Particle objects.  This
    particle list must be normalized.

    Returns: A resampled list of particles.  These will be NEW
    particles, the particles in the provided list will not be modified

    """
    #UNFINISHED

    # Implementation tips:

    # The probability that a particle is replicated should be
    # proportional to its weight.  The naive way of selecting a
    # particle is to randomly generate a number between 0 and 1, then
    # iterate through the particles accumulating their weights.  When
    # the accumulated total exceeds the random number, select that
    # particle.  Unfortunately, this is an O(n^2) algorithm: each
    # particle selection requires a (partial) pass through the
    # particle list.
    #
    # Alternatively, the algorithm from Table 4.4 in Thrun
    # et. al. Probalistic Robotics accomplishes the same thing in O(n)
    # time by selecting particles during a single pass through the
    # particle list.  Feel free to use either approach.

    return particle_list


def print_distribution(all_particles):
    for room in range(Particle.NUM_ROOMS):
        print("{:.3f}".format(calc_probability(all_particles, room)))
    print()

def run_demo():

    # Generate particles according to initial distribution:
    particles_room0 = [Particle(1.0, 0) for _ in range(7500)]
    particles_room1 = [Particle(1.0, 1) for _ in range(2500)]
    all_particles = particles_room0 + particles_room1

    normalize_particles(all_particles)

    # Predict...
    for p in all_particles:
        p.move(1) # (1 means right)

    print("B-(X1)")
    print_distribution(all_particles)

    # Correct...
    for p in all_particles:
        p.sense(1)  # (sensor says room 1)

    normalize_particles(all_particles)

    print("B(X1) (Before resampling)")
    print_distribution(all_particles)

    # Resample...
    all_particles = resample(all_particles)
    print("B(X1) (After resampling)")
    print_distribution(all_particles)


if __name__ == "__main__":
    run_demo()