% Model solution to task 4 of week 11
clear all; close all;

[mass,grad,nodes,element,interior,centers,ncoord,nelem] = FEMdata(3);

% Set up the diffusion coefficient
s = 100; % stochastic dimension
decay = 2.0; % decay of the input random field
indices = (1:s)';
deterministic = indices.^(-decay) .* sin(pi*indices*centers(:,1)') .* sin(pi* indices * centers(:,2)');
a = @(y) exp(deterministic'*y);

% Set up the loading
rhs = mass(interior,:)*nodes(:,1);

one = ones(length(interior),1);
gen = load('offtheshelf2048.txt');
z = gen(1:s);

R = 4;
nlist = 2.^(10:15);
rms = [];

for n = nlist
    disp(['n = ',num2str(n)]);
    results = [];
    for r = 1:R
        shift = rand(s,1);
        tmp = 0;
        parfor ii = 1:n
            qmcnode = norminv(mod(ii*z/n+shift,1));
            A = UpdateStiffness(grad,a(qmcnode));
            sol = A(interior,interior)\rhs;
            tmp = tmp + one'*mass(interior,interior)*sol;
        end
        results = [results;tmp/n];
    end
    qmcavg = mean(results);
    rmserror = norm(qmcavg-results)/sqrt(R*(R-1));
    rms = [rms;rmserror];
end

% Least squares fit
nadjusted = R*nlist;
lsq = [ones(size(nadjusted,2),1),log(nadjusted')]\log(rms);
lsq(1) = exp(lsq(1));

% Visualize the results
figure;
loglog(nadjusted,rms,'r.','MarkerSize',18), hold on
loglog(nadjusted,lsq(1)*nadjusted.^lsq(2),'b','LineWidth',2), hold off
set(gca,'FontSize',14,'FontWeight','bold')
legend('QMC error',['slope: ',num2str(lsq(2))]);
title(['QMC error (s = ',num2str(s),')'])
xlabel('n');
ylabel('error')