In [ ]:
from vpython import *
import numpy

scene = canvas()

#Classes
# A class is a template for an object
    #Ray class is a template for a ray
    #init: Set all of the variables for the beam
    #define methods

class ray():
    '''A light ray class'''
    def __init__(self,startPoint,directionVector,color):
        self.startPoint = startPoint
        self.currentPoint = startPoint
        self.directionVector = directionVector.norm()
        self.color = color
        self.length = 10
        self.ray = curve(pos=self.currentPoint, color=self.color)
        
        self.draw_ray()
        
        #Draw all rays in Source
    def draw_ray(self):
        '''Draw a single light ray'''
        newPoint = self.currentPoint + self.length*self.directionVector
        self.currentPoint = newPoint
        self.ray.append(newPoint)
        
        if newPoint.x > 5:
            self.currentPoint.x = 5
            thetaIncident = pi-diff_angle(self.directionVector,normal.directionVector)
            thetaRefracted = asin((n1/n2)*sin(thetaIncident))
            a1 = self.directionVector - normal.directionVector * (dot(self.directionVector,normal.directionVector))
            
            self.directionVector = (-1*cos(thetaRefracted)*normal.directionVector) + (sin(thetaRefracted)*a1.norm())
            
            newPoint = self.currentPoint + self.length*self.directionVector
            self.currentPoint = newPoint
            self.ray.append(newPoint)
            
        #Loop to check the position of newPoint.  If position is greater than 5, call Snell's Law
        #                                         and assign a new direction vector
    
        #Define normal as normal vector to object
        #Using Snell's Law (numpy.arcsin((1/(1.5))*sin(vector.diff_angle(ray1,normal))))
    def newDirection(self,newDirection):
        '''Set a new direction for the ray'''
        self.directionVector = newDirection                   
        
        #Stop when encounters an object and then reset length to 10
    def newLength(self,newLength):
        '''Set a new length for the vector'''
        self.length=newLength
        
class beam():
    '''A Beam Class'''
    def __init__(self, centerRay, width):
        self.centerRay = centerRay
        self.currentPoint = centerRay.currentPoint
        self.beamDirection = centerRay.directionVector
        self.width = width
        self.color = centerRay.color
        self.beam = []
        #PROBLEM HERE WITH Z-AXIS
        self.perp1 = vec(centerRay.directionVector.y,-centerRay.directionVector.x,0)
        self.perp2 = vec(centerRay.directionVector.cross(self.perp1))
    
        self.draw_beam()

    def draw_beam(self):
        '''Draw a beam; lots of parallel rays'''
        for i in numpy.arange(0,self.width,.1): #CAN WE DEFINE A BEAM USING A FOR LOOP?
            for k in numpy.arange(0,2*pi+0.1,.1):
                self.newStart = centerRay.startPoint + i*(self.perp1*cos(k)+self.perp2*sin(k))
            
                beamRay = ray(self.newStart, self.beamDirection, self.color)
                        
                self.beam.append(beamRay)
        
class pointSource():
    def __init__(self, position):
        self.position = position
        self.color = centerRay.color
        self.pointSource = []
        self.source = sphere(pos=vec(centerRay.startPoint), radius=0.01, color=self.color)
        
        self.draw_pointSource()
    
    def draw_pointSource(self):
        '''Draw a point source'''
        self.source
        
        for i in numpy.arange(0,2*pi+0.1,(pi/36)):
            for k in numpy.arange(0,2*pi+0.1,(pi/36)):
                
                x = centerRay.length*(cos(k)*cos(i))
                y = centerRay.length*(cos(k)*sin(i))
                z = centerRay.length*(sin(k))
                
                pointSourceRay = ray(self.position, vec(x,y,z), self.color)
                
                self.pointSource.append(pointSourceRay)
                

                
                

#####
# Start of World
#####

class world():
    def __init__(self):
        self.rays=[]
        self.objects=[]
        self.dL = 0.1   #Ray step size
        self.n = 1   #Default index of refraction for the world
        self.MAX_LENGTH = 100
        
        
    def add_ray(self,new_ray):
        '''Add a ray to the world'''
        self.rays.append(new_ray)
        new_ray.length=self.dL
        
    def add_object(self,new_object):
        '''Add a new object to the world'''
        self.objects.append(new_object)
        
    def draw_rays(self):
        #Loop for ceching Boundaries
        for i in self.rays:
            for j in self.objects:
                j.check_boundaries(i):
                i.draw_ray()
            
    #def draw_objects(self):
        
            
    def check_boundaries(self): #MOVE TO OBJECTS CLASS
        '''Determine if the position of the ray has crossed a boundary and if it has
        change direction of ray using Snells law.  Must run through everything in self.objects and see if the ray
        crossed one of the boundaries'''
        #Boundaries
        for i in self.rays:
            for j self.objects:
            b = i.currentPoint - j.position