% Master script that runs the whole thing
clear all;

% If moving is set to 1 then the camera 
% was moving and we have to compute the 
% table bounding box and geometry for
% every frame
moving = 0;

% Load table and ball image
table = imread('table.jpg');
ball = imread('ball.jpg');

[bX, bY] = size(ball);
halfBall = bX / 2;

imPrefix = 'Pic00';
imSuffix = '.jpg';

startPic = 20;
endPic = 130;

% Define the true table dimensions
% these will be used as reference points for
% the computation of the matrix that gives the 
% inverse perspective

% The table dimensions in cm
[tableHeight, tableWidth, o] = size(table);

% Define rotation matrix for the balls
R = [-1 0 0; 0 -1 0; 0 0 1];

for i = startPic : endPic
    
    disp(['Processing image ' int2str(i)]);
    % Read image
    
    % For the first video
    if i < 100
       imageFileName = [imPrefix '0' int2str(i) imSuffix];
    else
       imageFileName = [imPrefix int2str(i) imSuffix];
    end
    
    image = imread(imageFileName);    
    
    % Call the findLines function for the first image only
    % or for every frame if the camera is moving
    if i == startPic | moving == 1
       
        disp('Computing the table bounding box ...');
        
        % Get the edge detected version
        imGray = rgb2gray(image);
        imGray = double(imGray) / 255;  
        Edge = edge(imGray, 'canny', [.04 .07], 2); 

        [lines, points] = findLines(Edge);
        
        % Now that we have the corners of the table
        % on the uv plane we can compute the reference points
        % on the xy plane
        
        % Get the centroid of the corners
        cX = mean(points(:, 1));
        cY = mean(points(:, 2));
        
        % Now set the reference points so that they have the same
        % centroid. This must be done because we do not have any
        % information about the translation of the centers
        
        xyPoints = zeros(2, 4);
        
        xyPoints(1, 1) = cX - tableWidth / 2;
        xyPoints(2, 1) = cY - tableHeight / 2;
        
        xyPoints(1, 2) = cX + tableWidth / 2;
        xyPoints(2, 2) = cY - tableHeight / 2;
        
        xyPoints(1, 3) = cX + tableWidth / 2;
        xyPoints(2, 3) = cY + tableHeight / 2;

        xyPoints(1, 4) = cX - tableWidth / 2;
        xyPoints(2, 4) = cY + tableHeight / 2;

        % Get the inverse perspective matrix
        M = getPerspective(xyPoints, points');
        
        disp('Done');
        
    end
    
    disp('Detecting the balls ...');
    
    % Detect balls and return their centers on the image plane
    [my, mx] = findCircles(image, points, i);
    %drawLines(lines);
    disp('Done ...');
    
    % Get how many balls we have
    numBalls = length(my);
    one = ones(1, numBalls);
    
    disp('Drawing the pool table ...');
    ballCenters = [mx'; my'; one];
    temp = M * ballCenters;   
    
    % Homogenize the ball centers
    for k = 1 : numBalls
        temp(1 : 2, k) = temp(1 : 2, k) / temp(3, k);
    end
        
    newBallCenters = temp(1 : 2, :);
    
    corners = [points(:, 1)'; points(:, 2)'; [1 1 1 1]];
    newCorners = M * corners;
    
    for k = 1 : 4
        newCorners(1 : 2, k) = newCorners(1 : 2, k) / newCorners(3, k);
    end
        
    newCorners = newCorners(1 : 2, :);
    
    downX = max(min(newCorners, [], 2));
    downY = min(min(newCorners, [], 1));
    
    % Translate everything to 0, 0
    newCorners(1, :) = newCorners(1, :) - downX;
    newCorners(2, :) = newCorners(2, :) - downY;
    
    % Take the balls to the 0, 0 origin
    % The 18's are for the tables surrounding frame
    newBallCenters(1, :) = newBallCenters(1, :) - downX - 18;
    newBallCenters(2, :) = newBallCenters(2, :) - downY - 18;
    maxY = max(max(newCorners, [], 1));
    
    % Paint balls
    thisTable = table;
    for j = 1 : numBalls       
       for sr = 0 : 7
          for sth = 0 : ((10 - sr) / 35) : 2 * pi;
             thisTable(round(newBallCenters(2, j) + sr * cos(sth)), ...
                round(newBallCenters(1, j) + sr * sin(sth)), :) = 240 - 4 * sr ^ 2;
          end
       end
   end
   
   figure('color', 'w');
   imshow(thisTable);
   imwrite(thisTable, ['table' int2str(i) '.jpg']);
    
   disp('Done');
    
   % Let's free some memory
   clear image, imGray;
   disp('Hit any key to process next image ...');
   pause;
   close all;
        
end
