class Particle:
    def __init__(self):
        self.x = random.uniform(MIN_POSITION, MAX_POSITION)
        self.y = random.uniform(MIN_POSITION, MAX_POSITION)

        self.v = random.gauss(MEAN_VELOCITY, SIG_VELOCITY)
        theta = random.uniform(0, 2*np.pi)
        self.vx = np.cos(theta)
        self.vy = np.sin(theta)

        self.recent_collision = False

        #print("Initialized Particle: ")
        #print("x: ", self.x, " y: ", self.y, " v: ", self.v, " theta: ", self.theta)
        return

    def propagate(self, dt):
        n_x = self.x + self.vx*dt
        self.y = self.y + self.vy*dt


        if n_x < WALL_BOUNDARY < self.x:
            self.vx = -self.vx
            self.x = WALL_BOUNDARY
        elif self.x < WALL_BOUNDARY < n_x:
            if self.v < SPEED_BOUNDARY:
                self.vx = -self.vx
                self.x = WALL_BOUNDARY
            else:
                self.x = n_x
        else:
            self.x = n_x

        if self.x > MAX_POSITION or self.x < MIN_POSITION:
            self.vx = -self.vx
            if self.x > MAX_POSITION:
                self.x = MAX_POSITION
            else:
                self.x = MIN_POSITION
        if self.y > MAX_POSITION or self.y < MIN_POSITION:
            self.vy = -self.vy
            if self.y > MAX_POSITION:
                self.y = MAX_POSITION
            else:
                self.y = MIN_POSITION

        
        return



                      

      class AnimateParticles(object):
    def __init__(self, numparticles=PARTICLE_NUM):

        # create particles
        self.particles = [Particle() for i in range(numparticles)]

        # set up time delta to pass to stream
        dt = 0.05
        t = np.arange(0, 20, dt)

        self.stream = self.data_stream(dt)

        # velocities will be constant, so we can pull them out here
        self.v = [p.v for p in self.particles]
        
        # setup animated figure
        self.fig_setup()
        
        # then animate!
        self.ani = animation.FuncAnimation(self.fig, self.animate, len(t), init_func=self.ani_init, blit=True)


               
    def fig_setup(self):
        self.fig = plt.figure(figsize=[10,10], constrained_layout=True)
        gs = self.fig.add_gridspec(3,3)
        
        self.ax1 = self.fig.add_subplot(211, autoscale_on=False, xlim=(MIN_POSITION, MAX_POSITION), ylim=(MIN_POSITION, MAX_POSITION))
        #self.ax1.set_aspect('equal')
        self.ax1.axvline(WALL_BOUNDARY)
        self.ax1.set_xticks([])
        self.ax1.set_yticks([])

        self.ax2 = self.fig.add_subplot(629)
        self.ax2.set_visible(False)
        self.left_text = "# Particles on LEFT of boundary: %s"
        self.left_count = self.ax2.text(0.1, 1, "")


        self.ax3 = self.fig.add_subplot(6, 2, 10)
        self.ax3.set_visible(False)
        self.right_text = "# Particles on RIGHT of boundary: %s"
        self.right_count = self.ax3.text(0.1, 1, "")

        self.ax4 = self.fig.add_subplot(615)
        self.ax4.set_visible(False)
        
        self.ax_color = self.fig.add_subplot(915)
        self.ax_color.set_visible(False)

        return

    def ani_init(self):
        print("before figure iterate")
        x, y, _ = next(self.stream)

        self.scatter = self.ax1.scatter(x, y, marker='o', c=self.v, norm=cm.colors.Normalize(), cmap="YlOrRd")
        cbar = self.fig.colorbar(self.scatter, ax=self.ax_color, orientation="horizontal")
        cbar.set_label("speed")

        #self.ax2.hist(self.v)

        self.left_count.set_text('')
        self.right_count.set_text('')

        self.ax4.text(0.1, 1, "mean speed of particles: %s" % MEAN_VELOCITY)
        self.description = self.ax4.text(0.1, 0, "particle speed (gaussian):  mu %s  sig %s \n Maxwell demon speed cutoff: %s" % (MEAN_VELOCITY, SIG_VELOCITY, SPEED_BOUNDARY))
        
        return self.scatter, self.left_count, self.right_count, self.description

    def data_stream(self, dt):
        while True:
            [particle.propagate(dt) for particle in self.particles]
            points = [(p.x, p.y) for p in self.particles]
            x = [p.x for p in self.particles]
            y = [p.y for p in self.particles]
            yield x, y, points

    def animate(self, i):
        x, y, points  = next(self.stream)
        
        # set a and y on scatter plot
        #self.scatter = self.ax.scatter(x, y, c=self.v, cmap="rainbow")
        #self.scatter.set_offsets([0.01, 0.01])

        #self.plot.set_data(x, y)
        #self.text.set_text("animating")
        #self.scatter.set_offsets([(p.x, p.y) for p in self.particles])
        self.scatter.set_offsets(points)

        lcount = len([i for i in x if i < WALL_BOUNDARY])
        rcount = len([i for i in x if i > WALL_BOUNDARY])
        self.left_count.set_text(self.left_text % lcount)
        self.right_count.set_text(self.right_text % rcount)
        
        #return self.scatter,
        return self.scatter, self.left_count, self.right_count, self.description

    def save_animation(self ,path):
        Writer = animation.writers['ffmpeg']
        writer = Writer(fps=15)
        self.ani.save(path, writer=writer)