% [cdata,balls,energy] = step_function(cdata, balls)
% This function completes one iteration for the entire physical system of
% balls. It moves all balls with their respective velocities, manage
% collisions, and applies gravity. The input is the color matrix data
% (its size is required, not the content), and the set of balls
% (represented as a structure array). It returns the same and in addition
% the computed kinetic energy for the system.
function [cdata,balls,energy] = step_function(cdata,balls) %#codegen

%   Copyright 2010 The MathWorks, Inc.

% For every ball...
for i = 1:numel(balls)
    % Remember its old position
    oldBall = balls(i);
    % Move the ball based its own velocity
    balls(i) = move_ball(balls(i));
    % Rebounce if the ball is attempting to move outside the window
    [outside,ox,oy] = is_outside(cdata, balls(i));
    if outside
        balls(i) = oldBall;
        if ox
            balls(i).vx = -balls(i).vx;
        end
        if oy
            balls(i).vy = -balls(i).vy + 0.5; % 0.5 to apply gravity
        end
        % Reduce energy (for each collision)
        balls(i) = reduce_energy(balls(i));
        continue
    end
    % Does this ball hit some other ball?
    for j = 1:numel(balls)
        if i ~= j && does_collide(balls(i), balls(j))
            % Manage collision and reduce energy for both collided balls
            balls(i) = oldBall;
            balls = manage_collision(balls,i,j);
            balls(i) = reduce_energy(balls(i));
            balls(j) = reduce_energy(balls(j));
            break
        end
    end
end
% Compute the total kinetic energy of the system
energy = 0;
for i = 1:numel(balls)
    energy = energy + balls(i).vx*balls(i).vx + balls(i).vy*balls(i).vy;
end

% Sub-function to reduce energy; make the ball move slower
function ball = reduce_energy(ball)
energyLossFactor = 0.9;
ball.vx = ball.vx*energyLossFactor;
if ball.vx > 0
    ball.vx = ball.vx - 0.1;
end
if ball.vx < 0
    ball.vx = ball.vx + 0.1;
end
ball.vy = ball.vy*energyLossFactor;
if ball.vy > 0
    ball.vy = ball.vy - 0.1;
end
if ball.vy < 0
    ball.vy = ball.vy + 0.1;
end

% Sub-function to move ball and update velocity w.r.t. gravity
function ball = move_ball(ball)
ball.x = ball.x + ball.vx;
ball.y = ball.y + ball.vy;
ball.vy = ball.vy + 0.5;

% Sub-function to check whether two balls collide. We use the
% distance and radius of the balls to check whether the distance
% is smaller than the combined radius.
function y = does_collide(ball1, ball2)
dx = ball2.x - ball1.x;
dy = ball2.y - ball1.y;
rsum = ball1.r + ball2.r;
r2 = dx*dx + dy*dy;
y = r2 < rsum*rsum;

% balls = manage_collision(balls,i,j)
% Sub-function to manage collision of two balls indexed i and j.
function balls = manage_collision(balls,i,j)
% Never inline this function in the generated code. Makes the
% code more readable when doing code generation.
coder.inline('never');
% Extract x, y, vx, and vy properties from the balls.
xi = balls(i).x;
yi = balls(i).y;
xj = balls(j).x;
yj = balls(j).y;
vxi = balls(i).vx;
vyi = balls(i).vy;
vxj = balls(j).vx;
vyj = balls(j).vy;
% Compute normal of collision point
nx = xi-xj; 
ny = yi-yj;
n2=nx*nx+ny*ny;
n1=sqrt(n2);
nx = nx/n1;
ny = ny/n1;
% Compute velocities along the normal
vDotNi = vxi*nx+vyi*ny;
vDotNj = vxj*nx+vyj*ny;
% Compute the new velocities for the two collided balls
new_vxi = 2*vDotNj*nx - vxj;
new_vyi = 2*vDotNj*ny - vyj;
new_vxj = 2*vDotNi*nx - vxi;
new_vyj = 2*vDotNi*ny - vyi;
% If result is finite, update the structure array to reflect the collision
if isfinite(new_vxi) && isfinite(new_vyi) && isfinite(new_vxj) && isfinite(new_vyj)
    balls(i).vx = new_vxi;
    balls(i).vy = new_vyi;
    balls(j).vx = new_vxj;
    balls(j).vy = new_vyj;
end