function snooker(Vx,Vy)
% Function to animate a simple pool shot
% Vx, Vy are components of initial velocity
% of the cue ball.

close all

vmag = sqrt(Vx^2+Vy^2);
Vy = Vy/vmag;
Vx = Vx/vmag;

n_disks = 3;
disk_radius = 0.4;
box_size_x = 20.0;
box_size_y = 10;
pocket_width = 1.40;
restitution_wall = 0.8;
restitution_disk = 0.9;

% This sets the colors of the balls 0 is white, 1 is red, 2 is black
ball_color(1) = 0;
ball_color(2) = 2;
ball_color(3) = 1;

edges = set_pocket_edges; % Compute positions of edges of pockets

disk_wall_index = 1;    %specifies the disk that is closest to a wall
colliding_pair(1) = 1;  %specifies first disk in pair of closest disks
colliding_pair(2) = 1;  %specifies second disk in pair of closest disks
disk_edge_pair(1) = 1;  %specifies disk closest to an edge
disk_edge_pair(2) = 1;  %specifies edge closest to a disk
potted_ball = 1;         %specifies number of a potted ball

tstart = 0;
tstop = 66;

L = 20*disk_radius;
offset_y = 7*disk_radius;
offset_x = 10*disk_radius;
h2 = 7*disk_radius;
y0 = [-box_size_x/2+offset_x+L,-box_size_y/2+h2,Vx,Vy,
    -box_size_x/2+offset_x,-box_size_y/2+offset_y,0.0,0.,
    -box_size_x/2+L/2+offset_x,-box_size_y/2+h2/2+offset_y/2,0.,0.];
y0 = transpose(y0);

% We have to split the total simulation time into blocks otherwise
% MATLAB misses some collisions. This is a dirty fix for a pesky problem.
dtime = (tstop-tstart)/20;

    drawtable
    for n = 1:n_disks
    if (ball_color(n)==0)
       plot(y0(4*(n-1)+1),y0(4*(n-1)+2),'wo','MarkerSize',11,'MarkerFaceColor','white');
    elseif (ball_color(n)==1)
       plot(y0(4*(n-1)+1),y0(4*(n-1)+2),'ro','MarkerSize',11,'MarkerFaceColor','r');
    elseif (ball_color(n) == 2)
       plot(y0(4*(n-1)+1),y0(4*(n-1)+2),'ko','MarkerSize',11,'MarkerFaceColor','black');
    end
    end
    axis([-box_size_x/2-pocket_width,box_size_x/2+pocket_width,-box_size_y/2-pocket_width,box_size_y/2+pocket_width]);
    pause(0.5)


while(tstart<tstop)
Abstol = 10^(-08)*ones(4*n_disks,1);
options = odeset('RelTol',0.00000001,'AbsTol',Abstol,'Events',@detect_collision);

if (tstop-tstart>0.1)
  times = [tstart:0.1:tstart+dtime];
else
  times = [tstart,tstop];
end
[t,y,tevent,yevent,ievent] = ode113(@eom,times,y0,options);

for i=1:length(t)-1
    drawtable
    for n = 1:n_disks
    if (ball_color(n)==0)
       plot(y(i,4*(n-1)+1),y(i,4*(n-1)+2),'wo','MarkerSize',11,'MarkerFaceColor','white');
    elseif (ball_color(n)==1)
       plot(y(i,4*(n-1)+1),y(i,4*(n-1)+2),'ro','MarkerSize',11,'MarkerFaceColor','r');
    elseif (ball_color(n) == 2)
       plot(y(i,4*(n-1)+1),y(i,4*(n-1)+2),'ko','MarkerSize',11,'MarkerFaceColor','black');
    end
    end
    axis([-box_size_x/2-pocket_width,box_size_x/2+pocket_width,-box_size_y/2-pocket_width,box_size_y/2+pocket_width]);
    pause(0.025)
end
% Reset initial conditions for the next step
if (~isempty(tevent))
    y0 = collide(tevent,yevent,ievent);
else
    y0 = y(length(t),:);
end
tstart = t(length(t));
end

function dydt = eom(t,y)
% Equations of motion for the disks
    for nn=1:n_disks
        i = 4*(nn-1)+1;
        dydt(i) = y(i+2);
        dydt(i+1) = y(i+3);
        dydt(i+2) = 0;
        dydt(i+3) = 0.;
    end
    dydt = transpose(dydt);
end

function [eventvalue,stopthecalc,eventdirection] = detect_collision(t,y)
% Function to detect a collision of a ball with a cushion, the edge
% of a pocket, or another ball, and to detect a ball going into a pocket.
%
%  Find the shortest distance between a ball and a cushion (correcting for
%  pockets)
   dmin = box_size_x+box_size_y;
   for nn=1:n_disks
      i = 4*(nn-1)+1;
      d1 = box_size_x/2-(y(i)+disk_radius);
      d2 = (y(i)-disk_radius)+box_size_x/2;
      [d,min_wall] = min([d1,d2]);
      if (d<dmin)
%     Check for corner pockets
          z1 = y(i+1)+box_size_y/2-pocket_width/2;
          z2 = box_size_y/2-pocket_width/2-y(i+1);
          if (z1>0 && z2>0)
            dmin = d;
             disk_wall_index = nn;
          end
      end
      d3 = box_size_y/2-(y(i+1)+disk_radius);
      d4 = (y(i+1)-disk_radius)+box_size_y/2;
      [d,min_wall] = min([dmin,d3,d4]);
      if (d<dmin)
%     Check for corner pockets and middle pocket
          z1 = y(i)+box_size_x/2-pocket_width/2;
          z2 = box_size_x/2-pocket_width/2-y(i);
          z3 = abs(y(i))-pocket_width/2;
          if (z1>0 && z2>0 && z3>0)
            dmin = d;
             disk_wall_index = nn;
          end
      end
   end
%  First event is collision of disk with a wall
   eventvalue(1) = dmin;
   stopthecalc(1) = 1;
   eventdirection(1)=-1;
   dmin = (box_size_x+box_size_y);
%  Find the min dist between any pair of disks
   if (n_disks>1)
   for nn=1:n_disks-1
       i = 4*(nn-1)+1;
       for m = nn+1:n_disks
           j = 4*(m-1)+1;
           d = sqrt( (y(j)-y(i))^2 + (y(j+1)-y(i+1))^2)-2*disk_radius;
           if (d<dmin)
              dmin = d;
              colliding_pair(1)=nn;
              colliding_pair(2)=m;
           end
       end
   end
%  Second event is collision of two disks 
   eventvalue(2) = dmin;
   stopthecalc(2) = 1;
   eventdirection(2)=-1;
   
   end

   dmin = (box_size_x+box_size_y);
   for nn=1:n_disks
       i = 4*(nn-1)+1;
       for m = 1:12
           d = sqrt( (y(i)-edges(1,m))^2 + (y(i+1)-edges(2,m))^2)-disk_radius;
           if (d<dmin)
              dmin = d;
              disk_edge_pair(1)=nn;
              disk_edge_pair(2)=m;
           end
       end
   end
%  Third event is collision of disks with pocket edge 
   eventvalue(3) = dmin;
   stopthecalc(3) = 1;
   eventdirection(3)=-1;
%  Check for a ball leaving the box - if so it was potted
   dmin = box_size_x + box_size_y;
   for nn = 1:n_disks
      i = 4*(nn-1)+1;
      d1 = box_size_x/2-(y(i));
      d2 = (y(i))+box_size_x/2;
      d3 = box_size_y/2-(y(i+1));
      d4 = (y(i+1))+box_size_y/2;
      [d,min_wall] = min([d1,d2,d3,d4]);
      if (d<dmin)
         dmin = d;
         potted_ball = nn;
      end
   end
   %  Fourth event is a potted ball 
   eventvalue(4) = dmin;
   stopthecalc(4) = 1;
   eventdirection(4)=-1;
   
end

function y0 = collide(t,y,ievent)
%   Function to change velocity of disks when a collision occurs

y0 = y;
if (length(ievent)>1)
    y0=y(1,:);
end
for nn=1:length(ievent)
if (ievent(nn)==1)        %collision with a wall
    i = 4*(disk_wall_index-1)+1;
    d1 = box_size_x/2-(y0(i)+disk_radius);
    d2 = (y0(i)-disk_radius)+box_size_x/2;
    d3 = box_size_y/2-(y(i+1)+disk_radius);
    d4 = (y0(i+1)-disk_radius)+box_size_y/2;
    if (d1<0.0001 || d2<0.0001)   %collision with left or right wall
        y0(i+2) = -restitution_wall*y0(i+2);
    end
    if (d3<0.0001 || d4<0.0001) % collision with top or bottom wall
        y0(i+3) = -restitution_wall*y0(i+3);
    end
elseif (ievent(nn)==2)          %collision of two disks
    i = 4*(colliding_pair(1)-1)+1;
    j = 4*(colliding_pair(2)-1)+1;
    normal(1) = y0(j)-y0(i);
    normal(2)=y0(j+1)-y0(i+1);
    normal = normal/sqrt(dot(normal,normal));
    vn = (y0(j+2)-y0(i+2))*normal(1) + (y0(j+3)-y0(i+3))*normal(2);
    y0(i+2) = y0(i+2) + (1+restitution_disk)*vn*normal(1)/2;
    y0(i+3) = y0(i+3) + (1+restitution_disk)*vn*normal(2)/2;
    y0(j+2) = y0(j+2) - (1+restitution_disk)*vn*normal(1)/2;
    y0(j+3) = y0(j+3) - (1+restitution_disk)*vn*normal(2)/2;
elseif (ievent(nn)==3)          %collision of disk with pocket edge
    i = 4*(disk_edge_pair(1)-1)+1;
    j = disk_edge_pair(2);
    normal(1) = y(i)-edges(1,j);
    normal(2)= y(i+1)-edges(2,j);
    normal = normal/sqrt(dot(normal,normal));
    vn = y0(i+2)*normal(1) + y0(i+3)*normal(2);
    y0(i+2) = y0(i+2) - (1+restitution_wall)*vn*normal(1);
    y0(i+3) = y0(i+3) - (1+restitution_wall)*vn*normal(2);
elseif (ievent(nn)==4)          %potted ball - remove the ball that was potted
    for i = 1:4*(potted_ball-1)
         ynew(i) = y0(i);
     end
     for i = 4*potted_ball+1:4*n_disks
         ynew(i-4) = y0(i);
     end
     for i = 1:potted_ball-1
         bcnew(i) = ball_color(i);
     end
     for i = potted_ball+1:n_disks
         bcnew(i-1) = ball_color(i);
     end
     ball_color = bcnew;
     y0 = ynew;
     n_disks = n_disks-1;
end
end
end

function edges = set_pocket_edges
% Function to calculate coords of all pocket edges
    edges(1,1) = -box_size_x/2+pocket_width/2;
    edges(2,1) = -box_size_y/2;
    edges(1,2) = -box_size_x/2;
    edges(2,2) = -box_size_y/2 + pocket_width/2;
    edges(1,3) = box_size_x/2 - pocket_width/2;
    edges(2,3) = -box_size_y/2;
    edges(1,4) = box_size_x/2;
    edges(2,4) = -box_size_y/2 + pocket_width/2;
    edges(1,5) = box_size_x/2 - pocket_width/2;
    edges(2,5) = box_size_y/2;
    edges(1,6) = box_size_x/2;
    edges(2,6) = box_size_y/2 - pocket_width/2;
    edges(1,7) = -box_size_x/2+pocket_width/2;
    edges(2,7) = box_size_y/2;
    edges(1,8) = -box_size_x/2;
    edges(2,8) = box_size_y/2 - pocket_width/2;
    edges(1,9) = -pocket_width/2;
    edges(2,9) = -box_size_y/2;
    edges(1,10) = pocket_width/2;
    edges(2,10) = -box_size_y/2;
    edges(1,11) = -pocket_width/2;
    edges(2,11) = box_size_y/2;
    edges(1,12) = pocket_width/2;
    edges(2,12) = box_size_y/2;
end
function drawtable
    
    pocket(1,:) = [-box_size_x/2,-box_size_y/2];
    pocket(2,:) = [box_size_x/2,-box_size_y/2];
    pocket(3,:) = [box_size_x/2,box_size_y/2];
    pocket(4,:) = [-box_size_x/2,box_size_y/2];
    pocket(5,:) = [0,-box_size_y/2];
    pocket(6,:) = [0,box_size_y/2];
    cla
    hold on
    set(gca,'PlotBoxAspectRatio',[2,1,1]);
    plot(pocket(1:4,1),pocket(1:4,2),'-k','LineWidth',5);
    plot([-box_size_x/2,-box_size_x/2],[-box_size_y/2,box_size_y/2],'-k','LineWidth',5);
    plot(pocket(:,1),pocket(:,2),'bo','MarkerSize',15,'LineWidth',5,'MarkerFaceColor','white','MarkerEdgeColor','black');
    patch(pocket(1:4,1),pocket(1:4,2),[0,1,0]);
    plot(pocket(:,1),pocket(:,2),'bo','MarkerSize',15,'LineWidth',1,'MarkerFaceColor','white','MarkerEdgeColor','white');
end
end