from vpython import *
import numpy as np
import particle
import random
import math
import lennardjones as lj
import json

class box_sim():
    def __init__(self):
        self.delta_t = 0.0001 # incremental time
        self.t = 0
        self.box_thickness = 0.2
        self.particle_list = []
        self.size = 0
        self.s_size = 0.5
        self.sphere_list = []
        self.energy_label = None
        self.energy_value = 0.0
        self.mass = 1.0
        self.particle_number = 0
        self.epsilon = 1.0
        self.sigma = self.s_size * 2
        self.dict_data = {}
        self.dict_data['boxsim'] = []
        self.step = 0

    def draw_box(self, size):
        """Draws a cube with a transparent front

        :param: size, the size of the box
        """
        self.size = size
        bound = size / 2
        wallL = (box(pos=vec(-1.0 * bound - self.s_size, 0, 0),
                     size=vec(self.box_thickness, size + 1, size + 1),
                     color=color.green))
        wallR = (box(pos=vec(1.0 * bound + self.s_size, 0, 0),
                     size=vec(self.box_thickness, size + 1, size + 1),
                     color=color.green))
        wallU = (box(pos=vec(0, 1.0 * bound + self.s_size, 0),
                     size=vec(size + 1, self.box_thickness, size + 1),
                     color=color.blue))
        wallD = (box(pos=vec(0, -1.0 * bound - self.s_size, 0),
                     size=vec(size + 1, self.box_thickness, size + 1),
                     color=color.blue))
        wallB = (box(pos=vec(0, 0, -1.0 * bound - self.s_size),
                     size=vec(size + 1, size + 1, self.box_thickness),
                     color=color.red))

    def add_nov_particles(self, N):
        """Adds N particles with zero velocity

        :param: N number of particles
        """
        for i in np.arange(N):
            r_vi = np.zeros((3,))
            v_vi = np.zeros((3,))
            a_vi = np.zeros((3,))
            r_vi[0] = random.random() * self.size - (self.size / 2)
            r_vi[1] = random.random() * self.size - (self.size / 2)
            r_vi[2] = random.random() * self.size - (self.size / 2)
            particle_i = particle.particle(r_vi, v_vi, a_vi, i +
                                           self.particle_number)
            self.particle_list.append(particle_i)
        self.particle_number += N

    def add_cube_particles(self, nx, ny, nz):
        """Adds particles of dimension nx, ny, nz

        """
        N = nx * ny * nz
        for i in np.arange(nx):
            for j in np.arange(ny):
                for k in np.arange(nz):
                    r_vi = np.zeros((3,))
                    v_vi = np.zeros((3,))
                    a_vi = np.zeros((3,))
                    r_vi[0] = i * self.sigma * 2
                    r_vi[1] = j * self.sigma * 2
                    r_vi[2] = k * self.sigma * 2
        self.particle_number += N

    def add_two_particles(self):
        p1_r_vi = np.zeros((3,))
        p1_v_vi = np.zeros((3,))
        p1_a_vi = np.zeros((3,))
        p2_r_vi = np.zeros((3,))
        p2_v_vi = np.zeros((3,))
        p2_a_vi = np.zeros((3,))

        p2_r_vi[0] = self.sigma * 1.2

        particle1 = particle.particle(p1_r_vi,
                                      p1_v_vi,
                                      p1_a_vi,
                                      1 + self.particle_number)
        particle2 = particle.particle(p2_r_vi,
                                      p2_v_vi,
                                      p2_a_vi,
                                      2 + self.particle_number)
        self.particle_list.append(particle1)
        self.particle_list.append(particle2)
        self.particle_number += 2

    def add_single_particle(self, posx, posy, posz):
        r_vi = np.zeros((3,))
        v_vi = np.zeros((3,))
        a_vi = np.zeros((3,))

        r_vi[0] = posx * self.sigma
        r_vi[1] = posy * self.sigma
        r_vi[2] = posz * self.sigma

        s_particle = particle.particle(r_vi, v_vi, a_vi,
                                       1 + self.particle_number)
        self.particle_list.append(s_particle)
        self.particle_number += 1

    def add_v_particles(self, N, v):
        """Adds N particles with velocity v

        :param: N number of particles
        :param: v velocity
        """
        for i in np.arange(N):
            r_vi = np.zeros((3,))
            v_vi = np.zeros((3,))
            a_vi = np.zeros((3,))
            r_vi[0] = random.random() * self.size - (self.size / 2)
            r_vi[1] = random.random() * self.size - (self.size / 2)
            r_vi[2] = random.random() * self.size - (self.size / 2)
            vx = random.random()
            vy = random.random()
            vz = random.random()
            mag = math.sqrt(vx * vx + vy * vy + vz * vz)
            factor = v / mag
            v_vi[0] = factor * vx
            v_vi[1] = factor * vy
            v_vi[2] = factor * vz
            particle_i = particle.particle(r_vi, v_vi, a_vi, i +
                                           self.particle_number)
            self.particle_list.append(particle_i)
        self.particle_number += N

    def draw_particles(self):
        self.sphere_list = [None] * len(self.particle_list)
        for i, p in enumerate(self.particle_list):
            self.sphere_list[i] = (sphere(pos=vec(p.get_position()[0],
                                                  p.get_position()[1],
                                                  p.get_position()[2]),
                                    radius = self.s_size,
                                    color = color.cyan))

    def update_particles(self):
        for i, p in enumerate(self.particle_list):
            # apply lennard jones to get next acceleration
            next_a = self.particle_interactions(p)
            #next_a = np.zeros((3,))

            # check for collisions
            test_r = p.verlet_trialmove(self.delta_t)
            test_r += (self.size / 2)
            if test_r[0] > self.size or test_r[0] < 0.0:
                p.vel_flipx()
            if test_r[1] > self.size or test_r[1] < 0.0:
                p.vel_flipy()
            if test_r[2] > self.size or test_r[2] < 0.0:
                p.vel_flipz()

            # update position based off verlet algorithm
            p.verlet_move(self.delta_t)

            # update velocity based off verlet algorithm
            p.verlet_acc(next_a, self.delta_t)

            # update accleration
            p.update_acc(next_a)

            self.sphere_list[i].pos = vec(p.get_position()[0],
                                          p.get_position()[1],
                                          p.get_position()[2])
        self.t += self.delta_t
        self.step += 1

    def particle_interactions(self, particle_i):
        force_sum = np.zeros((3,))
        for j in self.particle_list:
            if particle_i.get_key() != j.get_key():
                # add to force_sum
                df = lj.get_force(particle_i, j, self.epsilon, self.sigma)
                force_sum = np.add(force_sum, df)
        acceleration = force_sum / self.mass
        return acceleration

    def init_numbers(self):
        self.energy_label = label(pos=vec((-0.8) * (self.size),
                                          (-0.8) * (self.size),
                                          0), text='%g'%self.energy_value)

    def calculate_energy(self):

        t_energy = 0.0
        tk_energy = 0.0
        tu_energy = 0.0
        for i, p in enumerate(self.particle_list):
            k_energy = 0.5 * self.mass * (p.get_velocity()[0] ** 2 +
                                          p.get_velocity()[1] ** 2 +
                                          p.get_velocity()[2] ** 2)
            u_energy = 0.0
            for j in self.particle_list:
                if p.get_key() != j.get_key():
                    du = lj.truncatedlennardjones(p, j, self.epsilon,
                                                  self.sigma)
                    u_energy += du

            tk_energy += k_energy
            tu_energy += u_energy
        t_energy = tk_energy + tu_energy
 
        if (self.step % 100) == 0:
            self.dict_data['boxsim'].append({
                'time': self.t,
                'kinetic_energy': tk_energy,
                'potential_energy': tu_energy,
                'total_energy': t_energy
            })

        self.energy_value = t_energy
        self.energy_label.text = '%g'%self.energy_value

    def write_data(self):
        with open('boxdata.json', 'w') as outfile:
            json.dump(self.dict_data, outfile)

if __name__ == '__main__':
    bs = box_sim()
    bs.draw_box(5)
    #bs.add_v_particles(100, 0.1)
    bs.add_two_particles()
    bs.add_single_particle(0, 1.2, 0)
    bs.add_single_particle(1.2, 1.2, 0.2)
    bs.draw_particles()
    bs.init_numbers()
    time_limit = 10
    while(bs.t < time_limit):
        rate(3000)
        bs.update_particles()
        bs.calculate_energy()
    bs.write_data()
