import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import time
import random as rn

class Application:
    def __init__(self):
        self.R = 1
        self.p1 = 1.5
        self.p2 = 1
        self.eta = 0.1
        self.L = 2
        self.to = time.time()
        self.t = time.time()

        self.scalex = 1
        self.scaley = 1
        self.scalez = 1
        self.fig = plt.figure(figsize=(10,5))
        self.ax = self.fig.add_subplot(111, projection='3d')
        plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

        self.plot = [self.ax.plot([0], [0], [0])]
        self.create_boundaries()
        self.scat = self.ax.scatter([],[],[], s=30)
        self.man = ParticuleManager(self.R, self.L, self.p2, self.p1, self.eta)

    def create_boundaries(self):
        i=0
        while i <= 2:
            self.plot[0] = self.ax.plot3D([0, self.L], [np.cos(np.pi*i), np.cos(np.pi*i)], [np.sin(np.pi*i), np.sin(np.pi*i)], c=[0,0,0,0.2])
            i = i + 0.25
        i = 0

        circle_angle = np.linspace(0, 2, 30)
        circle_y = np.linspace(0, 0, 30)
        circle_z = np.linspace(0, 0, 30)
        while i < 30:
            circle_y[i] = np.cos(circle_angle[i]*np.pi)
            circle_z[i] = np.sin(circle_angle[i]*np.pi)
            i = i + 1
        i = 0

        while i <= self.L:
            self.plot[0] = self.ax.plot3D(np.linspace(i, i, 30), circle_y, circle_z,  c=[0,0,0,0.2])
            i = i + 0.5
        self.ax.set_xlim(0, self.L)
        self.ax.set_ylim(-self.R, self.R)
        self.ax.set_zlim(-self.R, self.R)


    def update(self, i):
        t = time.time() - self.t
        self.t = time.time()
        self.man.move(t, self.p2, self.p1, self.eta, self.L, self.R)
        self.man.spawn(self.L)
        self.man.render(self.scat)

class Particule:
    def __init__(self, r=1, an=0, x=0, col=0, v=0):
        self.r = r
        self.an = rn.uniform(0,2)
        self.x = x
        self.y = r*np.cos(self.an*np.pi)
        self.z = r*np.sin(self.an*np.pi)
        self.v = v
        self.col = col

class ParticuleManager:
    def __init__(self, R, L, p2, p1, eta):
        self.time = time.time()
        self.nb = 700
        self.particules = []
        for i in range(self.nb):
            r=rn.uniform(0,R)
            x=rn.uniform(0,L)
            v=(p2-p1)*(r*r-R*R)/4/L/eta
            self.particules.append(Particule(r=r, x=x, v=v))
        #self.particules = [Particule(r=rn.uniform(0,R), x=rn.uniform(0,L)) for i in range(self.nb)]
        self.positionsx = ([0] * self.nb)
        self.positionsy = ([0] * self.nb)
        self.positionsz = ([0] * self.nb)

        for i in range(self.nb):
            self.positionsx[i] = self.particules[i].x
            self.positionsy[i] = self.particules[i].y
            self.positionsz[i] = self.particules[i].z

    def move(self, t, p2, p1, eta, L, R):
        i=0
        for p in self.particules:
            p.x = p.x + t*p.v
            self.positionsx[i]= p.x
            i = i+1

    def spawn(self, L):
        for i in range(self.nb):
            if self.particules[i].x > L:
                self.particules[i].x = 0
                self.positionsx[i] = self.particules[i].x


    def render(self, scat):
        scat._offsets3d = (self.positionsx, self.positionsy, self.positionsz)
a = Application()

ani = animation.FuncAnimation(a.fig, a.update, interval=40, blit=False)
plt.show()