% Sample Matlab code to test two-grid preconditioner for solving bundle adjustment problem

function test_two_grid

clc, clear

%% Data files dumped from SSBA code
%filename = 'JTJ49';
filename = 'JTJ138';
%filename = 'JTJ646';

addpath(genpath('/home/pawan/work/datasets')); % location to datasets

A = load(strcat(filename, '.txt'));
A = spconvert(A);
%A = Problem.A; 
A = A + triu(A,1)';
[m,n] = size(A);

%% number of eigenvectors to extract
nc2 = 3;
[Vl, ~] = eigs(A, nc2); % eigs gives the largest! We need largest!
P = Vl;

%% construct coarse grid
tic
Ac = P' * A * P;
toc

%% Do domain partitioning for Jacobi Smoother
nparts = 10;
[sizes_parts, ~, ~] = domain_decomposition(A, nparts);
nJ1 = sum(sizes_parts(1:nparts));
nJ2 = n - nJ1;

% Define Jacobi Smoother
J = blkdiag(A(1:nJ1, 1:nJ1),  A(nJ1+1 : n, nJ1+1: n));
[L, U] = lu(J);

%b = rand(m,1); 
tol = 1e-1; maxit = 500; im = 40;

x_known = rand(m,1);
b = A*x_known;

[~, flag, relres, its, ~] = gmres(A, b, im, tol, maxit, @tg_solve);

% get iterations of gmres
iter = (its(1)-1)*im + its(2);

fprintf('its = %d\nrelres = %g\nflag=%g\n', iter, relres, flag);

%% Two grid solver
function y = tg_solve(x)
    t = U \ (L \ x);
    g = P * (Ac \ (P' * x));
    q = A*g;
    y = t + g - U \ (L \ q);   
end

%% Function to do domain decomposition for Jacobi
function [sizes_parts, B,p] = domain_decomposition(A, nparts)
    p = size(A,1):-1:1;
    B = A;
    end_index = size(B,1);
    for i=size(B,1):-1:1
        if find(B(i,1:i),1) == find(B(1:i,i),1) && find(B(1:i,i),1) == i 
            end_index = i;
            break;
        end
    end
    for i=end_index:size(B,2)    
        if nnz(B(i,1:end_index-1)) ~= 0
            end_index = i-1;
            break;
        end
    end
    sizes_parts = [];
    ind = floor(end_index/nparts);
    prev = 1;
    for i=1:nparts-1
        while ~(find(B(ind,1:ind),1) == find(B(1:ind,ind),1) && find(B(1:ind,ind),1) == ind)
            ind = ind + 1;
        end
        sizes_parts = [sizes_parts, (ind-prev)];
        prev = ind;
        ind = ind + floor(end_index/nparts);
    end
    sizes_parts = [sizes_parts, (end_index-prev+1), (size(B,1) - end_index)];

end

