#!/usr/bin/env python2

#
# chaos - jzhu98
# See bottom for license.
#

from __future__ import division

from enum import Enum
from copy import deepcopy
from random import random as r
from visual import *
from octree import Octree


class Particle(sphere):
    def __init__(self, pos, vel, col, rad, mass):
        super(self.__class__, self).__init__(
            pos=pos,
            radius=rad,
            color=col
        )
        self.vel = vel
        self.mass = mass


class Dust(sphere):
    def __init__(self, pos, vel, col, rad, mass):
        super(self.__class__, self).__init__(
            pos=pos,
            radius=rad,
            color=col,
            make_trail=True
        )
        self.vel = vel
        self.mass = mass


def make_molecules(n_molecules, radius, mass):
    molecules = [Dust(
        pos=vector(0, 0, 0),
        vel=vector(0, 0, 0),
        col=color.red,
        rad=0.05,
        mass=1
    )]
    for i in xrange(n_molecules):
        molecules.append(Particle(
            pos=vector(r() - 0.5, r() - 0.5, r() - 0.5),
            vel=vector(r() - 0.5, r() - 0.5, r() - 0.5),
            col=color.blue,
            rad=radius,
            mass=mass
            ))
    return molecules


# Stable (no oscillating velocity) because return to inside bounding box
# is guaranteed
def calc_wall_collision(particle):
    box_bounds = 0.5
    if abs(particle.pos.x) >= box_bounds:
        particle.vel.x *= -1
    if abs(particle.pos.y) >= box_bounds:
        particle.vel.y *= -1
    if abs(particle.pos.z) >= box_bounds:
        particle.vel.z *= -1


def calc_part_collision(p, molecules):
    for m in molecules:
        # Avoid collision with itself
        if p is m:
            continue

        # If collision detected, perform elastic momentum transfer
        # (Python simultaneous assignment ftw!)
        if mag(p.pos - m.pos) <= p.radius + m.radius:
            # If masses are equal, velocities are swapped
            if abs(p.mass - m.mass) < 1e-3:
                p.vel, m.vel = m.vel, p.vel
            else:
                p.vel, m.vel = \
                    (p.vel * (p.mass - m.mass) + 2 * m.mass * m.vel) / (p.mass + m.mass), \
                    (m.vel * (m.mass - p.mass) + 2 * p.mass * p.vel) / (p.mass + m.mass)
            # Process only one collision per frame
            break


def update_velocity(molecules, dt):
    for m in molecules:
        calc_wall_collision(m)

    calc_part_collision(molecules[0], molecules)


def update_position(molecules, dt):
    for m in molecules:
        m.pos += m.vel * dt


def main():
    container = box(pos=vector(0, 0, 0),
                    size=vector(1, 1, 1),
                    color=color.blue,
                    opacity=0.1)

    # Generate list of 1 dust particle and 100 random small particles
    molecules = make_molecules(n_molecules=250, radius=0.01, mass=0.01)

    dt = 0.05

    while 1:
        rate(60)
        update_velocity(molecules, dt)
        update_position(molecules, dt)

if __name__ == '__main__':
    main()

# The MIT License (MIT)
# Copyright (c) 2014 James Zhu
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.