clear all;

%% Constants

% Physical constants
hbar = 1.052e-34;
q = 1.602e-19;

% Single-charge coupling energy (eV)
U_0 = 0.25;
% (eV)
kBT = 0.025;
% Contact coupling coefficients (eV)
gamma_1 = 0.005;
gamma_2 = gamma_1;
gamma_sum = gamma_1 + gamma_2;
% Capacitive gate coefficient
a_G = 0.5;
% Capacitive drain coefficient
a_D = 0.5;
a_S = 1 - a_G - a_D;

% Central energy level
mu = 0;

% Energy grid, from -1eV to 1eV
NE = 501;
E = linspace(-1, 1, NE);
dE = E(2) - E(1);
% TODO name this better
cal_E = 0.2;

% Lorentzian density of states, normalized so the integral is 1
D = (gamma_sum / (2*pi)) ./ ( (E-cal_E).^2 + (gamma_sum/2).^2 );
D = D ./ (dE*sum(D));

% Reference no. of electrons in channel
N_0 = 0;

voltages = linspace(0, 1, 101);

% Terminal Voltages
V_G = 0;
V_S = 0;

for n = 1:length(voltages)
    % Set varying drain voltage
    V_D = voltages(n);

    % Shifted energy levels of the contacts
    mu_1 = mu - V_S;
    mu_2 = mu - V_D;

    % Laplace potential, does not change as solution is found (eV)
    % q is factored out here, we are working in eV
    U_L = - (a_G*V_G) - (a_D*V_D) - (a_S*V_S);

    % Poisson potential must change, assume 0 initially (eV)
    U_P = 0;

    % Assume large rate of change
    dU_P = 1;

    % Run until we get close enough to the answer
    while dU_P > 1e-6
        % source Fermi function
        f_1 = 1 ./ (1 + exp((E + U_L + U_P - mu_1) ./ kBT));
        % drain Fermi function
        f_2 = 1 ./ (1 + exp((E + U_L + U_P - mu_2) ./ kBT));

        % Update channel electrons against potential
        N(n) = dE * sum( ((gamma_1/gamma_sum) .* f_1 + (gamma_2/gamma_sum) .* f_2) .* D);

        % Re-update Poisson portion of potential
        tmpU_P = U_0 * ( N(n) - N_0);
        dU_P = abs(U_P - tmpU_P);

        % Unsure why U_P is updated incrementally, perhaps to avoid oscillations?
        %U_P = tmpU_P;
        U_P = U_P + 0.1 * (tmpU_P - U_P)
    end

    % Calculate current based on solved potential.
    % Note: f1 is dependent on changes in U but has been updated prior in the loop
    I(n) = q * (q/hbar) * (gamma_1 * gamma_1 / gamma_sum) * dE * sum((f_1-f_2).*D);

end


%%Plotting commands

figure(1); 
h = plot(voltages, N,'k'); 
grid on;
set(h,'linewidth',[2.0]); 
set(gca,'Fontsize',[18]);
xlabel('Drain voltage [V]'); 
ylabel('Number of electrons');

figure(2); 
h = plot(voltages, I,'k'); 
grid on;
set(h,'linewidth',[2.0]); 
set(gca,'Fontsize',[18]);
xlabel('Drain voltage [V]'); 
ylabel('Current [A]');