%% Particle swarm optimization
%% The bees

x1 = [-5; 0];   %location bee 1
x2 = [-5; 0];   %location bee 2
x3 = [-5; 0];   %location bee 3
v1 = [0;0];     %velocity bee 1
v2 = [0;0];     %velocity bee 2
v3 = [0;0];     %velocity bee 3
%% The model

global phi1 phi2 phi3 omega
omega = .95;
phi1 = .2;
phi2 = .2;
phi3 = 0.2;
epsilon = .05;

p1 = x1; %best location bee 1 
p2 = x2; %best location bee 2
p3 = x3; %best location bee 3

MaxIt = 150;

locations = [x1,x2,x3]; %all explored locations

movement1 = repmat(x1,1);
movement2 = repmat(x2,1);
movement3 = repmat(x3,1);
pbest1 = gas_level(p1);
pbest2 = gas_level(p2);
pbest3 = gas_level(p3);
pg = best_location([p1,p2,p3]); %best overall location

cnt = 0;
while norm(pg) > epsilon && cnt <  MaxIt
    %animation
    plot(movement1(1,:),movement1(2,:), '.r', "MarkerSize", 15)
    hold on
    color = [gas_level(point),1-gas_level(point),0];
    plot(movement2(1,:),movement2(2,:), '.b', "MarkerSize", 15)
    plot(movement3(1,:),movement3(2,:), '.g', "MarkerSize", 15)
    plot(pg(1),pg(2), '.k', "MarkerSize", 15)  
    %plot(0,0, 'ko')
   
    for n = 1:length(locations)
        point = locations(:,n);
        if gas_level(point) > 0.6
            color = [1,0,0];
        elseif gas_level(point) > 0.2
            color = [1,0.5,0];
        else
            color = [0,1,0];
        end
        plot(point(1), point(2), '.' ,'color', color, 'MarkerSize', 10)
    end
    xlim([-5,5])
    ylim([-5,5])
    drawnow
    pause(.001)
    hold off
    xlim([-5, 5])
    ylim([-5, 5])
    
    %update velocities
    v1 = v_new(v1,x1,p1,pg);
    v2 = v_new(v2,x2,p2,pg);
    v3 = v_new(v3,x3,p3,pg);
    
    %update positions
    x1 = x_new(x1,v1);
    x2 = x_new(x2,v2);
    x3 = x_new(x3,v3);
    
    %add new explored locations 
    locations = [locations,x1,x2,x3];
    
    %update movement
    movement1 = [movement1(:,2:end), x1];
    movement2 = [movement2(:,2:end), x2];
    movement3 = [movement3(:,2:end), x3];
    
    %update best locations
    [p1,pbest1,p2,pbest2,p3,pbest3,pg] = new_best_locations(x1,x2,x3,p1,pbest1,p2,pbest2,p3,pbest3,pg);
    
    cnt = cnt + 1;
end

%% Plot color map

for n = 1:length(locations)
    point = locations(:,n);
    color = [gas_level(point),1-gas_level(point),0];
    plot(point(1), point(2), '.' ,'color', color, 'MarkerSize', 10)
    hold on
end
hold off
xlim([-5, 5])
ylim([-5, 5])
%% Functions
% Location comparison

function best_location = best_location(locations)
[~,I] = max(gas_level(locations));
best_location = locations(:,I);
end
% Update velocity

function v_new = v_new(v_old,x,p_i,p_g)
global phi1 phi2 phi3 omega
v_new = omega*v_old + phi1*rand*(p_i-x) + phi2*rand*(p_g-x) + phi3*(-1+2*rand(2,1));
end
% Update position

function x_new = x_new(x_old,v)
x_new = x_old + v;
end
% Update best locations

function [p1,pbest1,p2,pbest2,p3,pbest3,pg] = new_best_locations(x1,x2,x3,old_p1, old_pbest1,old_p2, old_pbest2,old_p3, old_pbest3,old_pg)
pg = old_pg;
if best_location([x1,old_p1]) == x1
    p1 = x1;
    pbest1 = gas_level(p1);
    if best_location([p1,pg]) == p1
        pg = p1;
    end
else 
     p1 = old_p1; 
     pbest1 = old_pbest1;
end

if best_location([x2,old_p2]) == x2
    p2 = x2;
    pbest2 = gas_level(p2);
    if best_location([p2,pg]) == p2
        pg = p2;
    end
else 
     p2 = old_p2; 
     pbest2 = old_pbest2;
end

if best_location([x3,old_p3]) == x3
    p3 = x3;
    pbest3 = gas_level(p3);
    if best_location([p3,pg]) == p3
        pg = p3;
    end
else 
     p3 = old_p3; 
     pbest3 = old_pbest3;
end
end
% Function with gas values

function gas_level = gas_level(x)
gas_level = exp(-(x(1,:).^2 + x(2,:).^2));
end