%============================================================
% Tony Hyun Kim
% 8.06: Hartree-Fock implementation on two hydrogen molecules
%       using four 1s-STO functions as the basis set.
%
% All files can be obtained at: 
%   http://web.mit.edu/kimt/www/8.06/paper/
%
% General notes:
%   - This is the second HF implementation I worked on for 
%       the 8.06 term paper. This has been written in a way that
%       corresponds exactly to the discussion given in the paper.
%   
% Dependencies:
%   sto1s.m: Defines the 1s Slater-type wavefunction
%============================================================
function [C eps phi] = hf2(r)

%============================================================
% Program parameters
%============================================================
dP = 1E-6;
maxIter = 500;
           
%============================================================
% Step 1: Enter the atomic coordinates and the basis functions
%============================================================
N = 4; % Number of electrons
M = 4; % Number of nuclei

Z = [1 1 1 1];   % Charges on nuclei

R = [0   0 0;... % Nuclear coordinates as a row 3-vector
     0.5 0 0;...
     r   0 0;...
     r+0.5 0 0];

% We will use Gaussians centered at each of the nuclei as 
% the basis set
n = 7;
for A = 1:M
    phi(n*(A-1)+1) = gaussian(1,R(A,:));
    phi(n*(A-1)+2) = gaussian(pi,R(A,:)+[1/pi 0 0]);
    phi(n*(A-1)+3) = gaussian(pi,R(A,:)-[1/pi 0 0]);
    phi(n*(A-1)+4) = gaussian(pi,R(A,:)+[0 1/pi 0]);
    phi(n*(A-1)+5) = gaussian(pi,R(A,:)-[0 1/pi 0]);
    phi(n*(A-1)+6) = gaussian(pi,R(A,:)+[0 0 1/pi]);
    phi(n*(A-1)+7) = gaussian(pi,R(A,:)-[0 0 1/pi]);
    %phi(n*(A-1)+8) = gaussian(1,R(A,:)+[2/pi 0 0]);
    %phi(n*(A-1)+9) = gaussian(1,R(A,:)-[2/pi 0 0]);
end
N_basis = n*M;

%============================================================
% STEP 2:
%============================================================
% The following matrices need to be computed only ONCE per a
% set of iterations.
%
% Compute S:  the overlap matrix,
%         T:  the kinetic energy matrix,
%         V_nucl: Nuclear potential
%         H_core:  core-Hamiltonian, H = T + V_nucl
%
%         I: The two-electron integrals
%============================================================

disp('Computing S,T,V_nucl...');
tic; % For benchmarking purposes. We can measure the amount
     % of time taken for Matlab to perform the computations
S      = zeros(N_basis,N_basis);
T      = zeros(N_basis,N_basis);
V_nucl = zeros(N_basis,N_basis);
for i = 1:N_basis
    for j = i:N_basis
        %------------------------------------------------------------
        % Overlap matrix, S
        %------------------------------------------------------------
        S(i,j) = gaussian_inner_product(phi(i),phi(j));
        S(j,i) = S(i,j);
        
        %------------------------------------------------------------
        % Kinetic energy matrix, T
        %------------------------------------------------------------
        T(i,j) = gaussian_kinetic_energy(phi(i),phi(j));
        T(j,i) = T(i,j);
        
        %------------------------------------------------------------
        % Nuclear potential energy, V_nucl
        %------------------------------------------------------------
        for A = 1:M % Sum over each of the nuclei
            V_nucl(i,j) = V_nucl(i,j) + ...
                          gaussian_nuclear_potential(phi(i),phi(j),...
                                                     Z(A),R(A,:));
        end
        V_nucl(j,i) = V_nucl(i,j);
    end
end
H_core = T + V_nucl;
time = toc;
msg = sprintf('  Done! (Time elapsed: %-.4f seconds)', time);
disp(msg);

%------------------------------------------------------------
% Now we compute the two electron integrals I
%------------------------------------------------------------
disp('Computing two-electron integrals...');
tic;
I = zeros(N_basis,N_basis,N_basis,N_basis);
for i = 1:N_basis
    for j = 1:N_basis
        for k = 1:N_basis
            for l = 1:N_basis
                I(i,j,k,l) = gaussian_two_electron(phi(i),phi(j),...
                                                   phi(k),phi(l));
            end
        end
    end
end
time = toc;
msg = sprintf('  Done! (Time elapsed: %-.4f seconds)', time);
disp(msg);

%============================================================
% STEP 3:
%------------------------------------------------------------
% - Obtain the transformation matrix X to be used later
%       (for symmetric orthogonalization)
%============================================================
X = zeros(N_basis,N_basis);
s_invroot = zeros(N_basis,N_basis);

% The function 'eig' gives the eigenvectors in columns of U. This is also
% the unitary matrix U s.t. U'*S*U is diagonal. (i.e. U'*S*U = s)
[U,s] = eig(S); 
for i = 1:N_basis
    s_invroot(i,i) = s(i,i)^(-1/2);
end
X = U*s_invroot*U';

%============================================================
% STEP 4:
%------------------------------------------------------------
% - Guess the density matrix P
%============================================================
P = zeros(N_basis,N_basis);

%============================================================
% STEP 5:
%------------------------------------------------------------
% We now ENTER THE SCF LOOP
%
% - Compute the matrix G from P and the two-electron integrals.
% - Compute the Fock matrix F = H_core + G;
%============================================================
disp('Entering SCF loop...');
F = zeros(N_basis,N_basis);

% Loop variables
count = 1;
sigma = 1;
clear disp;

% Enter SCF loop
while(sigma > dP)
    msg = sprintf('Iteration %d', count);
    disp(msg);
    
    P_prev = P;
    count = count + 1;
    
    %------------------------------------------------------------
    % Construct the matrix G
    %------------------------------------------------------------
    G = zeros(N_basis,N_basis);
    for i = 1:N_basis
        for j = 1:N_basis
            % BEGIN inner loop
            for p = 1:N_basis
                for q = 1:N_basis
                    G(i,j) = G(i,j) + P(p,q)*(I(i,q,j,p)-1/2*I(i,q,p,j));
                end
            end
            % END inner loop
        end
    end
    
    %============================================================
    % STEP 6:
    %------------------------------------------------------------
    % Construct the Fock matrix F
    %============================================================
    F = H_core + G;
    
    %============================================================
    % STEP 7:
    %------------------------------------------------------------
    % Obtain the transformed Fock matrix Fp = X'*F*X
    %============================================================
    Fp = X'*F*X;
    
    %============================================================
    % STEP 8:
    %------------------------------------------------------------
    % Diagonalize Fp to obtain C' and eps
    %============================================================
    [Cp, eps] = eig(Fp);
    
    %============================================================
    % STEP 9:
    %------------------------------------------------------------
    % Transform Cp to the C in F*C=S*C*eps
    %============================================================
    C = X*Cp;
    
    %============================================================
    % STEP 10:
    %------------------------------------------------------------
    % Form a new density matrix P
    %============================================================
    P = zeros(N_basis,N_basis);
    for i = 1:N_basis
        for j = 1:N_basis
            % BEGIN inner loop
            for k = 1:(N/2)
                P(i,j) = P(i,j) + 2*C(i,k)*C(j,k);
            end
            % END inner loop
        end
    end
    
    %============================================================
    % STEP 11:
    %------------------------------------------------------------
    % Determine whether the procedure has converged. We will monitor
    % the standard deviation of the successive density matrix
    % elements for this purpose.
    %============================================================
    sigma = 0;
    for i = 1:N_basis
        for j = 1:N_basis
            sigma = sigma + (P(i,j)-P_prev(i,j))^2;
        end
    end
    sigma = (1/N_basis^2*sigma)^(1/2);
    
    %============================================================
    % A brake mechanism. Force-quit if we exceed the specified
    %   maximum iteration.
    %============================================================
    if (count >= maxIter)
        break;
    end
end

%hf2_display(C(:,1),phi);