import numpy as np
from matplotlib.animation import FuncAnimation
from matplotlib import pyplot as plt

from mcl import Environment, MCL, Agent


def particle_weight(particle_observation, agent_observation):
    return 1 if particle_observation == agent_observation else 0


class Drawer(object):
    def __init__(self, ax,
                 range_, doors, initial_location, controls,
                 uncertaintity, covariance_ratio):
        environment = Environment(range_, doors, uncertaintity)
        self.mcl = MCL(range_, n_particles,
                       environment.observe, particle_weight,
                       covariance_ratio)
        self.agent = Agent(range_, initial_location, environment.observe)

        self.agent_scatter = ax.scatter([], [])

        self.particle_vlines = []
        for i in range(n_particles):
            vline = ax.axvline(ymin=0.2, ymax=0.8, color="cyan")
            self.particle_vlines.append(vline)

        ax.set_xlim(environment.range)
        ax.set_ylim([-0.2, 1.2])

        self.controls = controls

        xs = np.linspace(*range_, 200)
        doors = [1 if environment.is_door(x) else None for x in xs]
        self.door_lines = ax.plot(xs, doors, linewidth=6, color="r")

    def init(self):
        self.agent_scatter.set_offsets(np.c_[self.agent.location, 0])

        for x, vline in zip(self.mcl.particles, self.particle_vlines):
            vline.set_xdata(x)
        return self.particle_vlines + [self.agent_scatter]

    def __call__(self, i):
        control = self.controls[i]
        noise = np.random.normal(0, 0.3)

        self.agent.move(control + noise)

        particles = self.mcl.update_particles(control, self.agent.observation())

        self.agent_scatter.set_offsets(np.c_[self.agent.location, 0])

        for x, vline in zip(particles, self.particle_vlines):
            vline.set_xdata(x)
        return self.particle_vlines + [self.agent_scatter]


n_particles = 400
range_ = [-4, 24]
doors = [[-2, -1], [3, 5], [8, 12], [18, 22]]
initial_location = 2
uncertaintity = 0.05
covariance_ratio = 0.08

controls = [2, 5, -2, 1, 0, -3, -5, -2, 4, 2, 7, -3, -3, -4]
controls = np.vstack((controls, controls)).flatten()

fig, ax = plt.subplots()
drawer = Drawer(ax, range_, doors, initial_location, controls,
                uncertaintity, covariance_ratio)
animation = FuncAnimation(
    fig, drawer,
    init_func=drawer.init,
    frames=np.arange(len(controls)),
    interval=1000,
)
animation.save("mcl.mp4", dpi=400)