function [x1,v1,x2,v2]=chaos_balls(m1,m2,x1ini,v1ini,x2ini,v2ini)

% The function computes Poincare cross section for the problem of two 
% colliding balls in one dimension with gravity 

%PARAMETERS
dt=0.0001; %time step size (in seconds)
time_steps=2^22; %number of time steps (which defines duration of evolution) 
g=9.81; %gravitational acceleration

%VARIABLES
% m1,m2 - masses 
% x1ini,v1ini - initial position and velocity of the first ball
% x2ini,v2ini - position and velocity of the second ball

x1(1:time_steps)=0; %position of the first particle as a function of discrete time
v1(1:time_steps)=0; %position of the first particle as a function of discrete time 
x2(1:time_steps)=0; %position of the first particle as a function of discrete time
v2(1:time_steps)=0; %position of the first particle as a function of discrete time 
time(1:time_steps)=0; %real time
u1=0; %auxilary velocity variable
u2=0; %auxilary velocity variable
j=1; %counter for the number of points in Poincare section
% x1ps and v2ps are arrays, to be used for Poincare section plot, that will store only those values 
%of x2(i) and v2(i) when the collision between the balls occurs

%Initial conditions:
x1(1)=x1ini;
v1(1)=v1ini;
x2(1)=x2ini;
v2(1)=v2ini;

%Evolve differential equations using midpoint method and conservational of
%energy and momentum when the balls collide
for i=1:time_steps-1 %loop over discrete time steps
    
    time(i+1)=time(i)+dt;
    
    [x1(i+1),v1(i+1),x2(i+1),v2(i+1)]=update(x1(i),v1(i),x2(i),v2(i),dt);
    
    if (abs(x2(i+1)-x1(i+1))<eps || (x2(i+1)-x1(i+1))<0) %checks for collision between balls
   
       %??? update velocities here following conservation of energy and momentum formulae
       
       %??? after setting new velocities, call function update again to
       %recompute coordinates and velocity at i+1 step using new v1,v2
       %after the collision
        
       x2ps(j)=x2(i); v2ps(j)=v2(i); j=j+1; %Store data for Poincare section plot
     
    elseif (x1(i+1)<eps) %did first ball hit the floor?
    
        %??? new v1(i) velocity after the collistion
        x1(i)=abs(x1(i)); %this can be useful to prevent ball falling through the floor
        
        %??? after setting new velocity v1, call function update again to
        %recompute coordinates and velocity at i+1 step using new v1(i)
        %after the collision
        
    end %of the check for collision
       
    
end % loop over discrete time steps    

    figure(1), plot(x2ps,v2ps,'o');
    figure(2), plot(time,x2);  

end % function chaos_balls


function [x1new,v1new,x2new,v2new]=update(x1old,v1old,x2old,v2old,dt)
     
     g=9.81; %gravitational acceleration
    
     %below you should replace ??? with proper update according the the
     %midpoint method [in the order shown since formulas for, e.g., x1(i+1)
     %requires both v1(i) and v1(i+1)]
     
     v1new=???
     v2new=???
     x1new=???
     x2new=???

end %function update