% atoms = run_atoms(atoms,n)
% atoms = run_atoms(atoms,n,iter)
% Where 'atoms' the initial and final state of atoms (can be empty)
%       'n' is the number of atoms to simulate.
%       'iter' is the number of iterations for the simulation
%          (if omitted it is defaulted to 3000 iterations.)
function atoms = run_atoms(atoms,n,iter) %#codegen

%   Copyright 2010-2013 The MathWorks, Inc.

% Specify a designated structure name for this input variable
% to improve readability in the generated code.
coder.cstructname(atoms, 'Atom');

% We don't generate code for these extrinsic functions, instead we are using
% the existing MATLAB runtime environment. When we're doing code generation
% this functions will automatically be eliminated; they are only used for
% rendering.
coder.extrinsic('figure', 'set', 'image', 'gca', 'close', 'fprintf');

% Color data matrix
cdata = zeros(512,512,3,'uint8');

% Don't do rendering for code generation
if ~coder.target('rtw')
    fig = figure;
    pos = [0 0 0 0];
    pos = get(fig, 'Position');
    pos(2) = pos(2) - 64;
    pos(3) = 512;
    pos(4) = 512;
    set(fig, 'Position', pos);
    im = image(cdata);
    set(fig, 'Name', 'MATLAB Coder Atoms');
    set(gca, 'Visible', 'off');

    % If number of iterations is 0, then just render the data and return.
    if iter == 0
        render(im,cdata,atoms);
        return
    end
end

% Create an index matrix (using int16)
m = zeros(size(cdata,1),size(cdata,2),'int16'); 

% Initially the set of items is empty.
% x0 and dx0 is the "cursor" where the atoms come into the picture.
x0 = 16;
dx0 = 8;
for i = 1:iter
    % Add more atoms until we hit the upper limit (specified by 'n')
    if numel(atoms) < n
        % Add 16 atoms at a time, just to speed up the process
        for c = 1:16
            % Initial position for atom
            atom.x = x0;
            atom.y = 16;
            a = mod(numel(atoms),128);
            % Argument of velocity is changed
            arg = -pi*a/128;
            atom.vx = 3*cos(arg);
            atom.vy = 3*sin(arg);
            coder.cstructname(atom, 'Atom');
            % Put atom into list of atoms
            atoms = [atoms atom];
            % Move the "cursor" back and forth
            % (from where new atoms are added)
            x0 = x0 + dx0;
            if x0 > size(cdata,2) - 16
                x0 = size(cdata,2) - 16;
                dx0 = -dx0;
            end
            if x0 < 16
                x0 = 16;
                dx0 = -dx0;
            end
        end
    end
    % This is the core algorithm: Setup index matrix, run the step
    % function and render the result.
    [atoms,m] = setup_index_matrix(atoms,m);
    atoms = step_function(atoms, m);
    if ~strcmp(coder.target, 'rtw')
        % Skip the rendering if we generate a standalone application
        cdata = render(im,cdata,atoms);
        if mod(i,50) == 0
            fprintf(1, 'Iteration: %d\n', i);
        end
    end
end
if ~strcmp(coder.target, 'rtw')
    close(fig);
    fprintf(1, 'Completed iterations: %d\n', iter);
end

% render(im, cdata, atoms)
% Given a handle to a figure image (im), and a bitmap of color data (cdata)
% render the atoms.
function cdata = render(im, cdata, atoms)
coder.extrinsic('set', 'drawnow');
coder.inline('never');
cdata = zeros(size(cdata),class(cdata));
for i = 1:numel(atoms)
    x = round(atoms(i).x);
    y = round(atoms(i).y);
    cdata(y,x,:) = [255 255 255];
end
set(im, 'CData', cdata);
drawnow('update');

% [atoms,m] = setup_index_matrix(atoms,m)
% The matrix 'm' represents a discretized version of space.
% Each element of the matrix is set to the index of the atom occupying this
% space. We use all arguments as input and output to ensure that all values
% are passed by reference in the generated code.
function [atoms,m] = setup_index_matrix(atoms,m)
coder.inline('never');
m = zeros(size(m),class(m));
for i = 1:numel(atoms)
    x = round(atoms(i).x);
    y = round(atoms(i).y);
    if x >= 1 && x <= size(m,2) && y >= 1 && y <= size(m,1)
        m(y,x) = i;
    end
end 

% atoms = step_function(atoms,m)
% Move atoms based on their current position and velocities.
% Manage collisions if they are encountered.
function atoms = step_function(atoms,m)
% Never inline this function to improve readability in the generated code.
coder.inline('never');
for i = 1:numel(atoms)
    x = atoms(i).x;
    y = atoms(i).y;
    vx = atoms(i).vx;
    vy = atoms(i).vy;
    x = x + vx;
    y = y + vy;
    % Check for collision between atoms
    x = round(x);
    y = round(y);
    coll = false;
    if x >= 1 && x <= size(m,2) && y >= 1 && y <= size(m,1)
        j = m(y,x);
        if j > 0
            atoms = manage_collision(atoms, i, j);
            coll = true;
        end
    end
    atoms(i).x = atoms(i).x + atoms(i).vx;
    atoms(i).y = atoms(i).y + atoms(i).vy;
    if ~coll
        atoms(i).vy = atoms(i).vy + 0.1; % Gravity
    end
    % Check for walls and ground
    if atoms(i).x < 1
        atoms(i).vx = -atoms(i).vx;
        atoms(i).x = 1;
    end
    if atoms(i).x > size(m,2)
        atoms(i).vx = -atoms(i).vx;
        atoms(i).x = size(m,2);
    end
    if atoms(i).y < 1
        atoms(i).vy = -atoms(i).vy;
        atoms(i).y = 1;
    end
    if atoms(i).y > size(m,1)
        atoms(i).vy = -atoms(i).vy;
        atoms(i).vy = atoms(i).vy + 1; % Loose energy when hitting ground
        atoms(i).y = size(m,1);
    end
end

% atoms = manage_collision(atoms,i,j)
% Sub-function to manage collision of two atoms indexed i and j.
function atoms = manage_collision(atoms,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 atoms.
xi = atoms(i).x;
yi = atoms(i).y;
xj = atoms(j).x;
yj = atoms(j).y;
vxi = atoms(i).vx;
vyi = atoms(i).vy;
vxj = atoms(j).vx;
vyj = atoms(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 impulse along collision normal
vDotN = (vxi-vxj)*nx+(vyi-vyj)*ny;
Ix = vDotN*nx;
Iy = vDotN*ny;
% Compute the new velocities for the two collided atoms
new_vxi = vxi - Ix;
new_vyi = vyi - Iy;
new_vxj = vxj + Ix;
new_vyj = vyj + Iy;
% 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)
    atoms(i).vx = new_vxi;
    atoms(i).vy = new_vyi;
    atoms(j).vx = new_vxj;
    atoms(j).vy = new_vyj;
end
