clear all
close all

%%% This simulations aims to capture the group behaviors
%%% of fish in the presence of a shark-like predator

dt = 1; % timestep
tMax = 1000; % maximum time
L = 100; % length of domain
z = 1; % zoom on the interval
f = 0; % number of pictures saved (initialization)

dthreat = 40; % distance at which fish see sharks
maxSpeed = 2; % maximum speed
swimspeed = 2; % normal swimming speed
maxDcentroid = 35; % maximum distance fish are comfortable from centroid

schoolSize = 500; % number of fish per school
nSchools = 2; % number of school (1 or 2)
nSharks = 2; % number of sharks (1 or 2)
nTotal = nSchools*schoolSize + nSharks; 

% Initialize position and velocity arrays
xPosFish = zeros(1,nTotal);
yPosFish = zeros(1,nTotal);
zPosFish = zeros(1,nTotal);
xVelFish = zeros(1,nTotal);
yVelFish = zeros(1,nTotal);
zVelFish = zeros(1,nTotal);

% Set plot parameters
markersize = zeros(1,nTotal)+25;
colors = zeros(1,nTotal)+4;
colors(nSharks+1:schoolSize+nSharks) = 2;
colors(1:nSharks) = 5;
markersize(1:nSharks) = 1000;


%%% Set initial conditions
for j = nSharks+1:nSharks+schoolSize
    while sqrt((xPosFish(j)-L/2)^2+(yPosFish(j)-L/2)^2+(zPosFish(j))^2) > L/3
        xPosFish(j) = L*(rand());
        yPosFish(j) = L*(rand());
        zPosFish(j) = L*(rand()-0.5);
    end
    xVelFish(j) = - 1;
end
for j = nSharks+schoolSize+1:nSharks+nSchools*schoolSize
    while sqrt((xPosFish(j)+L/2)^2+(yPosFish(j)+L/2)^2+(zPosFish(j))^2) > L/3
        xPosFish(j) = -L*(rand());
        yPosFish(j) = -L*(rand());
        zPosFish(j) = L*(rand()-0.5);
    end
    xVelFish(j) = 1;
end


while f < 200 % Limit number of picture files
    
    % Begin time steps
    for t = 1:dt:tMax
        
        %%% Move the Sharks
        for i = 1:nSharks
            % Retrieve position and velocity values
            xPosShark = xPosFish(i);
            yPosShark = yPosFish(i);
            zPosShark = zPosFish(i);
            xVelShark = xVelFish(i);
            yVelShark = yVelFish(i);
            zVelShark = zVelFish(i);
            
            % Begin random swim
            r = rand();
            if r < 0.25
                xVelShark = xVelShark + 1;
                yVelShark = yVelShark + 1;
            elseif r < 0.5
                xVelShark = xVelShark - 1;
                yVelShark = yVelShark + 1;
            elseif r < 0.75
                xVelShark = xVelShark + 1;
                yVelShark = yVelShark - 1;
            else
                xVelShark = xVelShark - 1;
                yVelShark = yVelShark - 1;
            end
            
            % Swim away from 3L/4
            if abs(xPosShark) > 3*L/4
                xVelShark = -xVelShark;
            end
            if abs(yPosShark) > 3*L/4
                yVelShark = -yVelShark;
            end
            
            % Normalize velocities
            while sqrt(xVelShark^2+yVelShark^2) > maxSpeed+1
                xVelShark = 0.9*xVelShark;
                yVelShark = 0.9*yVelShark;
            end
            
            % Update shark positions and velocities
            xPosFish(i) = xPosShark + dt*xVelShark;
            yPosFish(i) = yPosShark + dt*yVelShark;
            zPosFish(i) = zPosShark + dt*zVelShark;
            xVelFish(i) = xVelShark;
            yVelFish(i) = yVelShark;
            zVelFish(i) = zVelShark;
        end
        
        
        
        %%% Move the Fish
        for j = 1:nSchools
            for i = nSharks+(j-1)*schoolSize+1:nSharks+j*schoolSize
                % Access position and velocity arrays
                xPosCur = xPosFish(i);
                yPosCur = yPosFish(i);
                zPosCur = zPosFish(i);
                xVelCur = xVelFish(i);
                yVelCur = yVelFish(i);
                zVelCur = zVelFish(i);
                
                % Initialize variable
                xtot = 0;
                ytot = 0;
                ztot = 0;
                xvtot = 0;
                yvtot = 0;
                zvtot = 0;
                
                % Calculate values needed for centroid
                for k = nSharks+(j-1)*schoolSize+1:nSharks+j*schoolSize
                    xtot = xtot + xPosFish(k);
                    ytot = ytot + yPosFish(k);
                    ztot = ztot + zPosFish(k);
                    xvtot = xvtot + xVelFish(k);
                    yvtot = yvtot + yVelFish(k);
                    zvtot = zvtot + zVelFish(k);
                end
                
                % Calculate centroid in three dimensions
                xcentroid = xtot/(schoolSize);
                ycentroid = ytot/(schoolSize);
                zcentroid = ztot/(schoolSize);
                dcentroid = sqrt((xcentroid-xPosCur)^2+(ycentroid-yPosCur)^2+(zcentroid-zPosCur)^2);
                
                % Calculate average velocities
                xVelAvg = xvtot/(schoolSize);
                yVelAvg = yvtot/(schoolSize);
                zVelAvg = zvtot/(schoolSize);
                
                % If too far from centroid, move towards centroid
                if dcentroid > maxDcentroid
                    dvxc = (xcentroid - xPosCur)/dcentroid;
                    dvyc = (ycentroid - yPosCur)/dcentroid;
                    dvzc = (zcentroid - zPosCur)/dcentroid;
                    
                    xVelCur = xVelCur + 0.5*swimspeed*dvxc;
                    yVelCur = yVelCur + 0.5*swimspeed*dvyc;
                    zVelCur = zVelCur + 0.5*swimspeed*dvzc;
                end
                
                % Rogue fish match velocities of the group
                if xVelCur > xVelAvg+1 || yVelCur > yVelAvg+1 || zVelCur > zVelAvg+1
                    xVelCur = xVelAvg;
                    yVelCur = yVelAvg;
                    zVelCur = zVelAvg;
                end
                
                % Drift around the origin
                if xPosCur > 0 && yPosCur > 0
                    xVelCur = xVelCur -0.1;
                elseif xPosCur < 0 && yPosCur > 0
                    yVelCur = yVelCur -0.1;
                elseif xPosCur < 0 && yPosCur < 0
                    xVelCur = xVelCur +0.1;
                elseif xPosCur > 0 && yPosCur < 0
                    yVelCur = yVelCur +0.1;
                end
                
                % Don't go too far from z = 0
                if zPosCur > L/4
                    zVelCur = zVelCur -0.2;
                elseif zPosCur < -L/4
                    zVelCur = zVelCur +0.2;
                end
                
                % Move away from close sharks
                for h = 1:nSharks
                    dShark = sqrt((xPosFish(h)-xPosCur)^2+(yPosFish(h)-yPosCur)^2+(zPosFish(h)-zPosCur)^2);
                    if dShark < dthreat
                        dvxShark = (xPosFish(h) - xPosCur)/dShark;
                        dvyShark = (yPosFish(h) - yPosCur)/dShark;
                        dvzShark = (zPosFish(h) - zPosCur)/dShark;
                        
                        xVelCur = xVelCur - 1.5*swimspeed*dvxShark;
                        yVelCur = yVelCur - 1.5*swimspeed*dvyShark;
                        zVelCur = zVelCur - 1.5*swimspeed*dvzShark;
                    end
                end
                
                % Normalize velocities
                while sqrt(xVelCur^2+yVelCur^2+zVelCur^2) > maxSpeed
                    xVelCur = 0.9*xVelCur;
                    yVelCur = 0.9*yVelCur;
                    zVelCur = 0.9*zVelCur;
                end
                
                % Update positions
                xPosFish(i) = xPosCur + dt*xVelCur;
                yPosFish(i) = yPosCur + dt*yVelCur;
                zPosFish(i) = zPosCur + dt*zVelCur;
                xVelFish(i) = xVelCur;
                yVelFish(i) = yVelCur;
                zVelFish(i) = zVelCur;
            end
            
            % Plot the positions of the fish and sharks
            scatter3(xPosFish,yPosFish,zPosFish,markersize,colors,'.');
            axis([-z*L, z*L, -z*L, z*L, -z*L, z*L]);
            view(45,45);
            pause(0.000001);
            
%             if mod(t+2,3) == 0
%                 f = f+1;
%                 saveas(gcf,['two_sharks_',sprintf('%03d',f),'.png']);
%             end
            
        end
    end
end




