% This code will simulate diffusion on a one dimensional domain using the
% auxiliary region method (Smith and Yates, 2017). It uses two auxiliary
% regions in order to couple a PDE and Brownian description of a reation
% diffusion system. For full details on implementation, please see the
% manuscript.
%
% Cameron Smith
% Date Created: 13/09/16
% Last Modified: 11/12/17

clear
close all

%% Initial Conditions 

% All particles in the left hand side of the domain
INITIAL_LEFT = 1;

% All particles in the right hand side of the domain
INITIAL_RIGHT = 0;

% Uniform initial condition
INITIAL_UNIFORM = 0;

% Domain empty for morphogen gradient
INITIAL_MORPH = 0;

% Test that only one has been chosen, and if not, display an error
INITIAL_SUM = INITIAL_LEFT + INITIAL_RIGHT + INITIAL_UNIFORM + INITIAL_MORPH;
if INITIAL_SUM ~= 1
    error('Please select one initial condition, there are currently %.0f selected',INITIAL_SUM)
end

% Create a variable that contains the initial condition as a string
if INITIAL_LEFT == 1
    INITIAL = 'LEFT';
elseif INITIAL_RIGHT == 1
    INITIAL = 'RIGHT';
elseif INITIAL_UNIFORM == 1
    INITIAL = 'UNIFORM';
elseif INITIAL_MORPH ==1
    INITIAL = 'MORPH';
end

%% Spatial Variables 

% Define the minimum and maximum values of the domain, and the dividing
% line between the two regimes
x_MIN = -1;
x_MAX = 1;
I = (x_MIN+x_MAX)/2;

% Define a PDE spatial step for the discretisation
h_PDE = 0.01;

% Create the PDE vector and find its length
x_VEC_PDE = x_MIN:h_PDE:I;
x_VEC_PDE_LEN = length(x_VEC_PDE);

% What width of bar should be used to disply the Brownian dynamics?
h_BRO = 0.1;

% Generate an error if this number doesn't divide the length of the
% Brownian domain
if mod((x_MAX-I)/h_BRO,1) ~= 0
    error('Please choose a different value for the bar width');
end

% The width of the auxiliary regions
WIDTH = 0.05;

% Pseudo-Compartment
PDE_AR = [I-WIDTH,I];

% Ghost Cell
BRO_AR = [I,I+WIDTH];

% Create a Brownian vector and find its length
x_VEC_BRO = I+0.5*h_BRO:h_BRO:x_MAX-0.5*h_BRO;
x_VEC_BRO_LEN = length(x_VEC_BRO);

% Create a vector that contains the Brownian domain split into PDE steps
% and a mesh with a much finer step size
x_VEC_PDE_BRO = I:h_PDE:x_MAX;
MESH = I+h_PDE/200:h_PDE/100:x_MAX-h_PDE/200;

% Finally create a vector that conatins the entire domain split into the
% PDE x_VEC_PDE_BRO
x_VEC_DOM = x_MIN:h_PDE:x_MAX;
x_VEC_DOM_LEN = length(x_VEC_DOM);

%% Temporal Variables 

% Define the minumum and maximum times
t_INIT = 0;
t_FINAL = 100; 

% Define a time-step
dt = 0.01;

% Define a PDE time step much smaller than dt
% dt_PDE = dt;

% Create a temporal vector and find its length
t_VEC = t_INIT:dt:t_FINAL;
t_VEC_LEN = length(t_VEC);

%% Model Constants and Functions 

% Number of repeats
M = 5000;

% Number of particles
N = 500;

% Diffusion coefficient
D = 0.0025;

% The microscopic jumping rate
d = D/WIDTH^2;

% Specify the production rate (morph gradient)
lambda = 400;

% Specify degredation rate (morph gradient)
mu  = 0.02;

% Change N if we have a morphogen gradient
if INITIAL_MORPH == 1
    N = lambda*D/mu;
end

% Value of r and s (morphogen gradient) in the theta-method scheme
r = D*dt/(h_PDE^2);
s = mu*dt;

% The value of theta for the theta method
% 0 = Explicit Euler
% 0.5 = Crank-Nicolson
% 1 = Implicit Euler
theta = 0.51;

% Generate the theta method matrix
if INITIAL_MORPH == 0
    C = Theta_C(theta,x_VEC_PDE_LEN,r);
else    
    % Calculate C
    [C,c] = Theta_C_Morph(theta,x_VEC_PDE_LEN,r,s,lambda,h_PDE,mu,D);
end

%% Recording Matrices 

% Record the PDE densities
SOL_PDE_MEAN = zeros(t_VEC_LEN,x_VEC_PDE_LEN);
SOL_PDE_VAR = zeros(t_VEC_LEN,x_VEC_PDE_LEN);

% Record the Brownian numbers
SOL_BRO_MEAN = zeros(t_VEC_LEN,x_VEC_BRO_LEN);
SOL_BRO_VAR = zeros(t_VEC_LEN,x_VEC_BRO_LEN);

% Record the number of particles in the PDE domain
NUM_PARTS_PDE_HYB = zeros(t_VEC_LEN,1);

% Record the number of particles in the Brownian domain
NUM_PARTS_BRO_HYB = zeros(t_VEC_LEN,1);

%% The Algorithm 

fprintf('Completing the algorithm:\n')

% The index in the PDE regime which is the lower end of the
% pseudo-compartment
PSEUD_IND = x_VEC_PDE_LEN - ((I-PDE_AR(1))/h_PDE);
if abs(round(PSEUD_IND)-PSEUD_IND) < 1e-10
    PSEUD_IND = round(PSEUD_IND);
else
    error('Invalid value of WIDTH')
end

% Loop through repeats
for m = 1:M
    
    % Give an indication as to where we are in the code for the console
    if M > 1
        if m == 1
            fprintf('Repetition number (out of %.0f): 1,',M)
        elseif m == M
            fprintf('%.0f...',M)
        else
            fprintf('%.0f,',m)
        end
    end
    
    % Initialise time
    t = t_INIT;
    
    % Pre-define some solution vectors
    PDE_SOL = zeros(1,x_VEC_PDE_LEN);
    BRO_SOL = zeros(1,x_VEC_BRO_LEN);
    if INITIAL_MORPH == 0
        a = zeros(1,2);
    else
        a = zeros(1,3);
    end
    
    % Using the initial condition, initialise the PDE solution vector
    if INITIAL_LEFT == 1
        PDE_SOL = N/((x_VEC_PDE_LEN-0.5)*h_PDE)*ones(1,x_VEC_PDE_LEN);
    elseif INITIAL_UNIFORM == 1 || INITIAL_MORPH == 1
        PDE_SOL = N/(2*(I-x_MIN))*ones(1,x_VEC_PDE_LEN);
    end
    
    % Count the number of particles in the PDE domain initially and find
    % how many should be in the Brownian part
    if INITIAL_LEFT == 1
        PARTICLES_BRO_HYB = 0;
    elseif INITIAL_RIGHT == 1
        PARTICLES_BRO_HYB = N;
    else
        PARTICLES_BRO_HYB = N/2;
    end
    
    PARTICLES_PDE_HYB = sum(PDE_SOL(2:end))*h_PDE;

    % Create a vector that contain the positions of all particles in the
    % Brownian side
    if INITIAL_LEFT == 0 || INITIAL_RIGHT == 0 || INITIAL_UNIFORM == 0
        n_BRO = rand(1,PARTICLES_BRO_HYB)*(x_MAX-I)+I;
    end
    
    % Calculate how many particles are in each bin provided we have
    % particles in the Brownian part
    if PARTICLES_BRO_HYB > 0
        BRO_SOL = histcounts(n_BRO,I:h_BRO:x_MAX);
    end
    
    % Add to solution matrices
    SOL_PDE_MEAN(1,:) = SOL_PDE_MEAN(1,:) + PDE_SOL;
    SOL_PDE_VAR(1,:) = SOL_PDE_VAR(1,:) + PDE_SOL.^2;
    SOL_BRO_MEAN(1,:) = SOL_BRO_MEAN(1,:) + BRO_SOL;
    SOL_BRO_VAR(1,:) = SOL_BRO_VAR(1,:) + BRO_SOL.^2;
    NUM_PARTS_PDE_HYB(1) = NUM_PARTS_PDE_HYB(1) + PARTICLES_PDE_HYB;
    NUM_PARTS_BRO_HYB(1) = NUM_PARTS_BRO_HYB(1) + PARTICLES_BRO_HYB;
    
    % Define the next PDE time step and its index
    NEXT_STEP_PDE = t + dt;
    NEXT_IND_PDE = 2;
    
    % Loop through time
    while t < t_FINAL
        
        % Before time
        t_BEFORE = t;
        
        % Count the number of particles in the pseudo-compartment
        NUM_PDE_AR = sum(PDE_SOL((PSEUD_IND+1):end))*h_PDE;
        
        % Create a logical that says if a particle is in the ghost cell or
        % not
        IN_BRO_AR = n_BRO < BRO_AR(2);
        
        % Count the number of particles in the ghost cell
        NUM_BRO_AR = sum(IN_BRO_AR);
        
        % Populate the propensity functions
        if NUM_PDE_AR < 1
            a(1) = 0;
        else
            a(1) = d*NUM_PDE_AR;  % Jump right from pseudo-compartment
        end
        a(2) = d*NUM_BRO_AR;  % Jump left from ghost cell
        if INITIAL_MORPH == 1
            a(3) = NUM_BRO_AR*mu;
        end
        
        % Sum up the propensities
        a0 = sum(a);
        
        % Calculate the time until next event
        tau = 1/a0*log(1/rand);
        
        % Create a temporary time
        t_TEMP = t+tau;
        
        % Decide if a particle jumps across the interface
        if t_TEMP < NEXT_STEP_PDE
            
            % Update time
            t = t_TEMP;
            
            % Generate a random number in (0,a0)
            a0_RAND = a0*rand;
            
            % Check which event we should enact
            % Enact a right jump
            if a0_RAND < a(1)
                
                % Remove a particle from PDE
                PDE_SOL(PSEUD_IND+1:end) = PDE_SOL(PSEUD_IND+1:end) - 1/WIDTH;
                
                % Add a particle to the Brownian domain
                n_BRO = [n_BRO,WIDTH*rand+I]; %#ok
                
                % Update the number of particles
                PARTICLES_BRO_HYB = PARTICLES_BRO_HYB + 1;
                PARTICLES_PDE_HYB = PARTICLES_PDE_HYB - 1;
                
            % Enact a left jump
            elseif a0_RAND < a(1) + a(2)
                
                % Add a particle to the PDE regime
                PDE_SOL(PSEUD_IND+1:end) = PDE_SOL(PSEUD_IND+1:end) + 1/WIDTH;
                
                % Remove a random particle in the ghost cell
                % Choose a random integer between 1 and NUM_GHOST
                PART_TO_MOVE = randi(NUM_BRO_AR);
                
                % Find that particles index in n_BRO
                k = 1;
                cumsuma = IN_BRO_AR(1);
                while cumsuma < PART_TO_MOVE
                    k = k+1;
                    cumsuma = cumsuma + IN_BRO_AR(k);
                end
                
                % Remove the particle
                n_BRO(k) = [];
                
                % Update the number of particles
                PARTICLES_BRO_HYB = PARTICLES_BRO_HYB - 1;
                PARTICLES_PDE_HYB = PARTICLES_PDE_HYB + 1;
                
            else
                % Remove a random particle from Brownian AR
                PART_TO_REMOVE = randi(NUM_BRO_AR);
                GHOST_SUM = cumsum(IN_BRO_AR);
                INDEX = 1;
                while GHOST_SUM(INDEX) < PART_TO_REMOVE
                    INDEX = INDEX + 1;
                end
                n_BRO(INDEX) = [];
                PARTICLES_BRO_HYB = PARTICLES_BRO_HYB - 1;                
            end
            
        else
            
            % Update time
            t = NEXT_STEP_PDE;
            
            % First update the Brownian and PDE regimes
            n_BRO = n_BRO + sqrt(2*D*dt)*randn(1,length(n_BRO));
            if INITIAL_MORPH == 0
                PDE_SOL = (C*PDE_SOL')';
            else
                PDE_SOL = (C*PDE_SOL')' + c';
                PARTICLES_PDE_HYB = sum(PDE_SOL(2:end))*h_PDE;
%                 PARTICLES_PDE_HYB = Trap_Rule(PDE_SOL,h_PDE);
%                 PARTICLES_PDE_HYB = Simpson(PDE_SOL,h_PDE);
            end
            
            % Enact the boundary conditions on the Brownian particles
            RIGHT_ESC = n_BRO > x_MAX;
            LEFT_ESC = n_BRO < I;
            n_BRO(RIGHT_ESC) = n_BRO(RIGHT_ESC) - 2*(n_BRO(RIGHT_ESC) - x_MAX);
            n_BRO(LEFT_ESC) = n_BRO(LEFT_ESC) + 2*(I - n_BRO(LEFT_ESC));
            
            % Find if a degredation reaction occurs
            if INITIAL_MORPH == 1
                REMOVE = mu*dt > rand(PARTICLES_BRO_HYB,1);
                REMOVE_PARTS = ((REMOVE + ~IN_BRO_AR') == 2);
                if sum(REMOVE) > 0
                    n_BRO(REMOVE_PARTS) = [];
                    PARTICLES_BRO_HYB = PARTICLES_BRO_HYB - sum(REMOVE_PARTS);
                end
            end
    
            % Calculate how many particles are in each bin provided we have
            % particles in the Brownian part
            if PARTICLES_BRO_HYB > 0
                BRO_SOL = histcounts(n_BRO,I:h_BRO:x_MAX);
            end
            
            % Record
            SOL_PDE_MEAN(NEXT_IND_PDE,:) = SOL_PDE_MEAN(NEXT_IND_PDE,:) + PDE_SOL;
            SOL_PDE_VAR(NEXT_IND_PDE,:) = SOL_PDE_VAR(NEXT_IND_PDE,:) + PDE_SOL.^2;
            SOL_BRO_MEAN(NEXT_IND_PDE,:) = SOL_BRO_MEAN(NEXT_IND_PDE,:) + BRO_SOL;
            SOL_BRO_VAR(NEXT_IND_PDE,:) = SOL_BRO_VAR(NEXT_IND_PDE,:) + BRO_SOL.^2;
            NUM_PARTS_PDE_HYB(NEXT_IND_PDE) = NUM_PARTS_PDE_HYB(NEXT_IND_PDE) + Trap_Rule(PDE_SOL,h_PDE);
            NUM_PARTS_BRO_HYB(NEXT_IND_PDE) = NUM_PARTS_BRO_HYB(NEXT_IND_PDE) + PARTICLES_BRO_HYB;
        
            % Update the next PDE step and its index
            if t ~= t_FINAL
                NEXT_IND_PDE = NEXT_IND_PDE + 1;
                NEXT_STEP_PDE = t_VEC(NEXT_IND_PDE);
            end
            
        end 
    end
end

fprintf('DONE\n\n')

% Average over all repeats
SOL_PDE_MEAN = SOL_PDE_MEAN/M;
SOL_PDE_VAR = SOL_PDE_VAR/M-(SOL_PDE_MEAN).^2;
SOL_BRO_MEAN = SOL_BRO_MEAN/M;
SOL_BRO_VAR = SOL_BRO_VAR/M-(SOL_BRO_MEAN).^2;
NUM_PARTS_PDE_HYB = NUM_PARTS_PDE_HYB/M;
NUM_PARTS_BRO_HYB = NUM_PARTS_BRO_HYB/M;

%% PDE Across Entire Domain 

% Pre-define a matrix containing the solution
MEAN_FIELD = zeros(t_VEC_LEN,x_VEC_DOM_LEN);

% Initial condition
% Using the initial condition, initialise the PDE solution vector
if INITIAL_LEFT == 1
    MEAN_FIELD(1,1:x_VEC_PDE_LEN) = N/((x_VEC_PDE_LEN-0.5)*h_PDE)*ones(1,x_VEC_PDE_LEN);
elseif INITIAL_RIGHT == 1
    MEAN_FIELD(1,(x_VEC_PDE_LEN+1):end) = N/((x_VEC_PDE_LEN-0.5)*h_PDE)*ones(1,x_VEC_DOM_LEN-x_VEC_PDE_LEN);
elseif INITIAL_UNIFORM == 1 || INITIAL_MORPH == 1
    MEAN_FIELD(1,:) = N/(2*(I-x_MIN))*ones(1,x_VEC_DOM_LEN);
end

% Recalculate C_MAT
if INITIAL_MORPH == 0   
    C = (Theta_C(theta,x_VEC_DOM_LEN,r));
else
    [C,c] = Theta_C_Morph(theta,x_VEC_DOM_LEN,r,s,lambda,h_PDE,mu,D);
end
% Loop through time
for j = 2:t_VEC_LEN
    
    % Progress the solution
    if INITIAL_MORPH == 0
        MEAN_FIELD(j,:) = (C*MEAN_FIELD(j-1,:)')';
    else
        MEAN_FIELD(j,:) = (C*MEAN_FIELD(j-1,:)')' + c';
    end
    
end

%% Saving Workspace

fprintf('Saving Workspace...')

% Name of the file
NAME = ['WORKSPACE_ARM_INITIAL_',INITIAL];

% Saving
save(NAME)

fprintf('DONE\n\n')