clear;close all;clc
rng(0);

load('MNIST.mat')
[w_tr, h_tr, l_tr] = size(training.images);

training.images = training.images(:,:,randperm(60000, 6000));

d = 784 - 1;
Manifold = spherefactory(d+1, 1);
Manifold.type = 'sphere';
Manifold.retr = @Manifold.exp;
Manifold.invretr = @Manifold.log;
Manifold.transp = @Manifold.isotransp;

nagents = 50;
lncostterms = ceil(6000/nagents);
%%% construct agents

prob = cell(nagents,1);
opts = cell(nagents,1);

delta_i = 1e-4;

Li = [];
sigma2 = [];
v = 1e-3;
for i = 1:nagents
    Z = [];
    ttt = training.images(:,:,(i-1)*lncostterms+1 : i*lncostterms);
    for j = 1:lncostterms
        ttttt  = ttt(:,:,j);
        Z = [Z, ttttt(:)];
    end
    Zdiags = diag(Z'*Z);
    Li(i) = 2 * mean(Zdiags);
    D = cell(lncostterms,1);
    for j = 1:lncostterms
        D{j} = sparse(Z(:,j));
    end

%     sigma2(i) = log(1/delta_i) * Li(i)^2 / lncostterms / lncostterms;
    sigma2(i) = log(1/delta_i) * 1.00^2 / lncostterms / lncostterms;

    probtmp.M = Manifold;
    probtmp.lncostterms = lncostterms;
    probtmp = agent_Prob_PCA(D, lncostterms, probtmp);
    prob{i} = probtmp;

    optstmp.tau0 = 1.;
    optstmp.tau1 = 1.;
    optstmp.tau = 1.;
    opts{i} = optstmp;       
end

L = sum(Li) / nagents;
Lbar = L;

problem.M = Manifold;
problem.nagents = nagents;
problem.agent = prob;
problem.cost = @(x) cost(problem, x);
problem.grad = @(x) grad(problem, x);
problem.partialgrad = @(x, S) partialgrad(problem, x, S);

options.agent = opts;

rng(1)
x0 = Manifold.rand();


% for step size
% options.stepsize_type = 'decay';
options.stepsize_type = 'fix';
% options.stepsize = 0.2 / Lbar;
options.stepsize = 0.001;
options.decaysteps = 10;
options.option = 2;
options.st = 1;

seed = 1;

% SGD
options.maxiter = 500;
options.localMethod = 'DPRSGD';
K = 4;
epsilon_i = 1000;
for i = 1:nagents
    options.agent{i}.lmaxiter = K;
    options.agent{i}.sigma = 0 * sqrt(K * sigma2(i) / epsilon_i / epsilon_i);
    options.agent{i}.batch_size = ceil(lncostterms*.999);
    options.agent{i}.stepsize = options.stepsize;
end
fvs1 = [];
ngfs1 = [];
for i = 1:seed
    [x1, fv1, info1] = PriRFed(problem, x0, options);
    fvs1 = [fvs1, info1.fvs];
    ngfs1 = [ngfs1, info1.num_gfs];
end
figure(1)
plot(mean(ngfs1, 2), mean(fvs1,2),'LineWidth',2);
leg = ["RFed-SGD"];

epsilon_i = 0.6;
for i = 1:nagents
    options.agent{i}.lmaxiter = K;
    options.agent{i}.sigma = sqrt(K * sigma2(i) / epsilon_i / epsilon_i);
    options.agent{i}.batch_size = ceil(lncostterms*0.999);
    options.agent{i}.stepsize = options.stepsize;
end
fvs2 = [];
ngfs2 = [];
for i = 1:seed
    [x2, fv2, info2] = PriRFed(problem, x0, options);
    fvs2 = [fvs2, info2.fvs];
    ngfs2 = [ngfs2, info2.num_gfs];
end
hold on
figure(1)
plot(mean(ngfs2, 2), mean(fvs2,2),'LineWidth',2);
leg = [leg,"PriRFed-SGD-" + string(epsilon_i)];

epsilon_i = 1.;
for i = 1:nagents
    options.agent{i}.lmaxiter = K;
    options.agent{i}.sigma = sqrt(K * sigma2(i) / epsilon_i / epsilon_i);
    options.agent{i}.batch_size = ceil(lncostterms*0.999);
    options.agent{i}.stepsize = options.stepsize;
end
fvs3 = [];
ngfs3 = [];
for i = 1:seed
    [x3, fv3, info3] = PriRFed(problem, x0, options);
    fvs3 = [fvs3, info3.fvs];
    ngfs3 = [ngfs3, info3.num_gfs];
end
hold on
figure(1)
plot(mean(ngfs3, 2), mean(fvs3,2),'LineWidth',2);
leg = [leg,"PriRFed-SGD-" + string(epsilon_i)];

epsilon_i = 1.4;
for i = 1:nagents
    options.agent{i}.lmaxiter = K;
    options.agent{i}.sigma = sqrt(K * sigma2(i) / epsilon_i / epsilon_i);
    options.agent{i}.batch_size = ceil(lncostterms*0.999);
    options.agent{i}.stepsize = options.stepsize;
end
fvs4 = [];
ngfs4 = [];
for i = 1:seed
    [x4, fv4, info4] = PriRFed(problem, x0, options);
    fvs4 = [fvs4, info4.fvs];
    ngfs4 = [ngfs4, info4.num_gfs];
end
hold on
figure(1)
plot(mean(ngfs4, 2), mean(fvs4,2),'LineWidth',2);
leg = [leg,"PriRFed-SGD-" + string(epsilon_i)];
ylabel('Function value, $f(x_t)$', 'interpreter','latex', 'FontSize', 18)
xlabel('#grad','FontSize',18)
legend(leg)
title('fixed stepsize')
saveas(gca, 'PCA_Syn_SGD_fixed_Knot1_cost.fig')

% SD
[x_sd, fv_sd, info_sd] = steepestdescent(problem, x0);
f_star = min([info_sd.cost]);

figure(2)
semilogy(mean(ngfs1,2), mean(fvs1,2) - f_star,...
    mean(ngfs2,2), mean(fvs2,2) - f_star,...
    mean(ngfs3,2), mean(fvs3,2) - f_star, ...
    mean(ngfs4,2), mean(fvs4,2) - f_star,'LineWidth',2)
xlabel('#grad','FontSize',18)
ylabel('Excess risk, $f(x_t) - f(x^{\star})$','interpreter','latex', 'FontSize', 18)
legend(leg)
title('fixed stepsize')

saveas(gca, 'PCA_Syn_SGD_fixed_Knot1_ER.fig')

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function out = cost(problem, x)
    out = 0;
    for i = 1:problem.nagents
        out = out + problem.agent{i}.cost(x);
    end
    out = out / problem.nagents;
end

function out = grad(problem, x)
    out = 0;
    for i = 1:problem.nagents
        out = out + problem.agent{i}.grad(x);
    end
    out = out / problem.nagents;
end

function out = partialgrad(problem, x, S)
    out = 0;
    for i = reshape(S, 1,length(S))
        out = out + problem.agent{i}.grad(x);
    end
    out = out / length(S);
end












