close all;clear;clc;
rng(0);

d = 2;
Manifold = sympositivedefinitefactory(d);
Manifold.type = 'SPD';
Manifold.retr = @Manifold.exp;
Manifold.invretr = @Manifold.log;
Manifold.transp = @Manifold.paralleltransp;

numSamples = 6000;

nagents = 40;
lncostterms = ceil(numSamples / nagents);

fprintf('numSamples:%d, nagents:%d, ncostterms:%d\n', numSamples, nagents, lncostterms)

t=0;
Z = cell(numSamples, 1);
while t < numSamples
    tmp = wishrnd(eye(d)/ 2, 2);
    if dist(tmp, eye(d)) <= 1
       Z{t+1} = tmp;
       t = t+1;
    end    
end


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

delta_i = 1e-4;
sigma2 = [];
for i = 1:nagents

    D = {Z{(i-1)*lncostterms+1 : i*lncostterms}};

    sigma2(i) = log(1/delta_i) * 0.1^2 / lncostterms / lncostterms;

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

    optstmp.tau0 = 0.1;
    optstmp.tau1 = 0.1;
    optstmp.tau = 0.1;
    optstmp.sigma = sqrt(sigma2(i));
    opts{i} = optstmp;       
end

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);

% checkgradient(problem)

options.agent = opts;

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


% for step size
% options.stepsize_type = 'decay';
options.stepsize_type = 'fix';
options.stepsize = 0.1;
options.decaysteps = 20;
options.option = 2;
options.st = 2;

seed = 5;

% SGD
options.maxiter = 100;
options.localMethod = 'DPRSGD';
K = 1;
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 = ["PriRFed-SGD-\infty"];

epsilon_i = 0.3;
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 = 0.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
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 = 0.5;
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,'FontSize',12)
title('fixed stepsize','FontSize',12)

saveas(gca, 'FM_Syn_SGD_fixed_K_1_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,'FontSize',12)
title('fixed stepsize','FontSize',12)

saveas(gca, 'FM_Syn_SGD_fixed_K_1_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


function out = dist(x, y)
    sm = sqrtm(x);
    tmp = real(logm(sm \ y / sm));
    out = sqrt( sum(sum(tmp .* tmp)) );
end






