% N SPHERES IN A TWO-LEVEL SYSTEM:

% Initial values:
N=100;
KE=3924;                 % initial kinetic energy in microJoules
PE=980;                  % Potential energy in microJoules
totaltimesteps=10;
timestep=1;              % Set time step to 1
L=100;                   % Platform width is 100 cm
r=0.1;                   % Sphere radius in cm
m=1;                     % Set particle mass in grams

% CREATE X-Y POSITION VECTOR:
x=L*rand(N,2);         % Spheres are randomly distributed on lower level 
                         % of side length L

% CREATE X-Y VELOCITY VECTOR
v=zeros(N,2);          % Zeros serve as temporary placeholders
vi=sqrt(2*KE/m);  % Solve for initial particle speed

% Initial velocity is vi in a random direction on the x-y plane.
% Create initial velocity vector entries as follows:
for i=1:N
    r1=rand;
    v(i,1)=vi*cos(2*pi*r1);
    v(i,2)=vi*sin(2*pi*r1);
end

% CREATE A VECTOR TO RECORD UPPER/LOWER LEVEL FOR EACH PARTICLE:
level=zeros(N,1);    % 0=lower level, 1=upper level; All start on lower (0)

% CREATE A VECTOR TO TRACK NUMBER OF PARTICLES IN UPPER LEVEL:
N1=zeros(totaltimesteps+1,1);  % First entry reflects initial count 
                                   % of 0 particles; later we will record 
                                   % additional entries
         
% MAKE THE PARTICLES MOVE
for i=1:totaltimesteps % At each time interval, perform the following 
                   % operations:

    for j=1:N  % Calculate new position/velocity for each particle:
        % UPDATE POSITION:
        x(j,1)=x(j,1)+timestep*v(j,1);
        x(j,2)=x(j,2)+timestep*v(j,2);
        
        % PERFORM POSITION CHECKS:
        %----------------------------------------------------------------%
        if x(j,1)<0                 % If particle has position x<0:
            x(j,1)= -x(j,1);      % reflect x-position over y-axis
            v(j,1)= -v(j,1);      % reflect x-component of velocity
        end
       
        %----------------------------------------------------------------%
        if x(j,1)>L                 % If particle has position x>L:
            x(j,1)=2*L-x(j,1);    % reflect x-position over x=L
            v(j,1)= -v(j,1);      % reflect x-component of velocity
        end
        
        %----------------------------------------------------------------%
        if x(j,2)<0                 % If particle has position y<0:
            x(j,2)= -x(j,2);      % reflect y-position over x-axis
            v(j,2)= -v(j,2);      % reflect y-component of velocity
        end
        
        %----------------------------------------------------------------%
        if x(j,2)>2*L               % If particle has position y>2L:
            x(j,2)=4*L-x(j,2);    % reflect y-position over y=2L
            v(j,2)= -v(j,2);      % reflect y-component of velocity
        end        
        
        %----------------------------------------------------------------%
        if level(j,1)==0 && x(j,2)>L % If sphere was previously on lower 
                                       % level and is now at position y>L,   
           KE = (v(j,1)*v(j,1)+v(j,2)*v(j,2))/2; %KE of particle j
           if KE > PE         % then, if its KE exceeds mgh,
              ratio=sqrt(1-PE/KE);     
                                       % calculate the ratio of new speed 
                                       % to old speed, i.e. v(new)/v(old).
                                       % (Note that ratio<1 in this case.)
                                       
              v(j,1)=ratio*v(j,1); % Then scale x- and y- velocity
              v(j,2)=ratio*v(j,2); % components accordingly.
           
                                       % Otherwise,
           else x(j,2)=2*L-x(j,2); % reflect position over y=L            
                v(j,2)= -v(j,2);   % reflect y-component of velocity                                   
           end
        end
        
        %----------------------------------------------------------------% 
        if level(j,1)==1 && x(j,2)<L % If sphere was previously on upper 
                                       % level and is now at position y<L,      
        KE = (v(j,1)*v(j,1)+v(j,2)*v(j,2))/2; %KE of particle j
        ratio=sqrt(1+PE/KE);           % Calculate the ratio of new speed 
                                       % to old speed, i.e. v(new)/v(old).
                                       % (Note that ratio>1 in this case.)
                                       
        v(j,1)=ratio*v(j,1);       % Then scale x- and y- velocity
        v(j,2)=ratio*v(j,2);       % components accordingly.
        end
    
        % END OF POSITION CHECKS      
        %----------------------------------------------------------------%
        
        % UPDATE LEVEL VECTOR WITH A 0 OR 1:
        if x(j,2)<L               % If sphere is below y=L,
             level(j,1)=0;          % it is on the lower level (0).
        else level(j,1)=1;          % Otherwise it is on the upper level(1).
        end
    end
        
    % CHECK FOR COLLISIONS:
    for j=1:N             % For each of the N spheres,
        for k=j+1:N       % check the distance to all remaining spheres.
            if (x(j,1)-x(k,1))^2 + (x(j,2)-x(k,2))^2 < (2*r)^2 ...
                    && level(j,1)+level(k,1)~=1 
                          % If the spheres are on the same level (i.e.
                          % the sum is 0+0 or 1+1) and the distance between
                          % centers is less than 2r, then they collide.
        
            % Find the unit vector from sphere j to sphere k:
            vectorjk=[x(j,1)-x(k,1)  x(j,2)-x(k,2)];
            unitjk=vectorjk/norm(vectorjk);
            
            % Velocity of sphere j in the "parallel" direction of unitjk:
            vjc=dot(v(j,1:2),unitjk)*unitjk;
            % Velocity of sphere j perpendicular to unitjk:
            vjperp=v(j,1:2)-vjc;
            
            % Velocity of sphere k in the "parallel" direction of unitjk:
            vkc=dot(v(k,1:2),unitjk)*unitjk;
            % Velocity of sphere k perpendicular to unitjk:
            vkperp=v(k,1:2)-vkc;
            
            % To obtain final velocities, swap "parallel" velocities 
            % as follows:
            v(j,1:2)=vkc+vjperp;
            v(k,1:2)=vjc+vkperp;  
            
            end          
        end
             
     end       
        
    % Record the number of spheres on the upper level at this time interval
    N1(i+1,1)=sum(level);
    
end

% Create a vector to record the number of spheres on the lower level
N0=N*ones(totaltimesteps+1,1)-N1;
