%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Simple Electrostatic Particle In Cell (PIC) code in MATLAB
%
% Herve Grabas, hgrabas@uchicago.edu
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%reset variables
clear variables

%identify globals needed by the potential solver
global EPS0 QE den A phi0

%setup constants
EPS0 = 8.854e-12;		%permittivity of free space
QE = 100*1.7622e-19;		    %elementary charge
K = 1.381e-23;			%boltzmann constant
M = 9.1093e-31;         %electron mass unit

%input settings
phi0 = 0;              %reference potential


%set simulation domain
nx = 20;               %number of nodes in x direction
ny = 20;               %number of nodes in y direction
ts = 100;              %number of time steps
dh = 0.001;            %cell size
nb_part = 1e4;        %number of particle
dt = 1e-12;	           %time step

%compute some other values
nn = nx*ny;             %total number of nodes
Lx = (nx-1)*dh;         %domain length in x direction
Ly = (ny-1)*dh;         %domain length in y direction

%allocate particle array
part_x = zeros(nb_part,2); %particle positions
part_v = zeros(nb_part,2); %particle velocities

%set up multiplication matrix for potential solver
%here we are setting up the Finite Difference stencil

A = zeros(nn);              %allocate empty nn * nn matrix

%set regular stencil on internal nodes
for j=2:ny-1                    %only internal nodes
    for i=2:nx-1
        u = (j-1)*nx+i;         %unknown (row index)
        
        A(u,u) = -4/(dh*dh);    %phi(i,j)
        A(u,u-1)=1/(dh*dh);     %phi(i-1,j)
        A(u,u+1)=1/(dh*dh);     %phi(i+1,j)
        A(u,u-nx)=1/(dh*dh);    %phi(i,j-1)
        A(u,u+nx)=1/(dh*dh);    %phi(i,j+1)
    end
end

%neumann boundary on y=0
for i=1:nx
    u=i;
    A(u,u) = -1/dh;              %phi(i,j)
    A(u,u+nx) = 1/dh;            %phi(i,j+1)
end

%neumann boundary on y=Ly
for i=1:nx
    u=(ny-1)*nx+i;
    A(u,u-nx) = 1/dh;            %phi(i,j-1)
    A(u,u) = -1/dh;              %phi(i,j)
end

%dirichlet boundary on x=Lx
for j=1:ny
    u=(j-1)*nx+nx;
    A(u,:)=zeros(1,nn);         %clear row
    A(u,u) = 1;                 %phi(i,j)
end

%dirichlet boundary on x=0
for j=1:ny
    u=(j-1)*nx+1;
    A(u,:)=zeros(1,nn);         %clear row
    A(u,u) = 1;                 %phi(i,j)
end

%initialize
phi = ones(nx,ny)*phi0;         %set initial potential to phi0
np = nb_part;                   %initialize number of particles


%insert particles randomly distributed in y and in the first cell
part_x(1:nb_part/2,1)=rand(nb_part/2,1)*dh+Lx/2+dh/2;       %x position
part_x(nb_part/2+1:nb_part,1)=rand(nb_part/2,1)*dh+Lx/2+5*dh/2;
part_x(:,2)=rand(nb_part,1)*dh+Ly/2-dh/2;      %y position

%sample Maxwellian in x and y, add drift velocity in x
part_v(:,1)=0;
part_v(:,2)=0;

disp(['Solving potential for the first time. Please be patient, this could take a while.']);

%%%%%%%%%%%%%%%%%%%%%%%%
% MAIN LOOP
%%%%%%%%%%%%%%%%%%%%%%%%
for it=1:ts                     %iterate for ts time steps
    
    %reset field quantities
    den = zeros(nx,ny);         %number density
    efx = zeros(nx,ny);         %electric field, x-component
    efy = zeros(nx,ny);         %electric field, y-component
    chg = zeros(nx,ny);         %charge distribution
    
    %*** 1. CALCULATE CHARGE DENSITY ***
    
    % deposit charge to nodes
    for p=1:np                          %loop over particles
        fi = 1+part_x(p,1)/dh;          %real i index of particle's cell
        i = floor(fi);                  %integral part
        hx = fi-i;                      %the remainder
        
        fj = 1+part_x(p,2)/dh;          %real j index of particle's cell
        j = floor(fj);                  %integral part
        hy = fj-j;                      %the remainder
        
        %interpolate charge to nodes
        chg(i,j) = chg(i,j) + (1-hx)*(1-hy);
        chg(i+1,j) = chg(i+1,j) + hx*(1-hy);
        chg(i,j+1) = chg(i,j+1) + (1-hx)*hy;
        chg(i+1,j+1) = chg(i+1,j+1) + hx*hy;
    end
    
    %calculate density
    den = chg/(dh*dh);
    
    %apply boundaries
    den(1,:) = 2*den(1,:);      %double density since only half volume contributing
    den(nx,:) = 2*den(nx,:);
    den(:,1) = 2*den(:,1);
    den(:,ny) = 2*den(:,ny);
    
    %*** 2. CALCULATE POTENTIAL ***
    %phi = eval_2dpot_GS(phi);
    phi = poicalc(phi,dh,dh,nx,ny);
    
    %*** 3. CALCULATE ELECTRIC FIELD ***
    efx(2:nx-1,:) = phi(1:nx-2,:) - phi(3:nx,:);  %central difference on internal nodes
    efy(:,2:ny-1) = phi(:,1:ny-2) - phi(:,3:ny);  %central difference on internal nodes
    efx(1,:) = 2*(phi(1,:) - phi(2,:));           %forward difference on x=0
    efx(nx,:) = 2*(phi(nx-1,:) - phi(nx,:));      %backward difference on x=Lx
    efy(:,1) = 2*(phi(:,1) - phi(:,2));           %forward difference on y=0
    efy(:,ny) = 2*(phi(:,ny-1) - phi(:,ny));      %forward difference on y=Ly
    
    efx = efx / (2*dh);    %divide by dominator
    efy = efy / (2*dh);
    
    
    %*** 5. MOVE PARTICLES ***
    p=1;
    while(p<=np)                        %loop over particles
        fi = 1+part_x(p)/dh;            %i index of particle's cell
        i  = floor(fi);
        hx = fi-i;                      %fractional x position in cell
        
        fj = 1+part_x(p,2)/dh;          %j index of particle' cell
        j = floor(fj);
        hy = fj-j;                      %fractional y position in cell
        
        %gather electric field
        E=[0 0];
        E = [efx(i,j) efy(i,j)]*(1-hx)*(1-hy);      %contribution from (i,j)
        E = E+ [efx(i+1,j) efy(i+1,j)]*hx*(1-hy);   %(i+1,j)
        E = E + [efx(i,j+1) efy(i+1,j)]*(1-hx)*hy;  %(i,j+1)
        E = E + [efx(i+1,j+1) efy(i+1,j+1)]*hx*hy;  %(i+1,j+1)
        
        %update velocity and position
        F = QE*E;                           %Lorentz force, F=qE
        a = F/M;                            %acceleration
        part_v(p,:) = part_v(p,:)+a*dt;     %update velocity
        part_x(p,:) = part_x(p,:)+part_v(p,:)*dt;   %update position
        
        %process boundaries
        
        %absorbing boundary on left, right, top or if in object
        if (part_x(p,1)<0 || part_x(p,1)>=Lx || part_x(p,2)>=Ly || part_x(p,2)<0)
            part_x(p,:) = part_x(np,:);         %kill particle by replacing it with last particle
            part_v(p,:) = part_v(np,:);
            np = np - 1;                        %reduce particle count
            p = p-1;                            %reduce particle index so this entry will be processed again
        end
        
        p=p+1;                                  %move to the next particle
    end
    
    %*** 6. PLOT RESULTS ***
    %if (mod(it,4) == 0 || it==ts)              %plot only every 25 time steps
    
    figure(1);                              %density
    contourf(den');
    title(sprintf('Density %i', it));
    colorbar;
    
    figure(2);                              %potential
    contourf(phi');
    title(sprintf('Potential %i', it));
    colorbar;
    
    figure(3);
    quiver(1+part_x(:,2)'/dh,1+part_x(:,1)'/dh,part_v(:,2)',part_v(:,1)');
    
    figure(5);
    scatter(part_x(:,1),part_x(:,2));
    
    drawnow;                                %force draw
    %end
    
    disp(sprintf('Time Step %i, Particles %i',it, np));
    
    if (np == 0)
        disp('No more particles -> exit');
        break;
    end
    
end

disp(sprintf('Complete!\n'));                   %exit