close all;clear;clc;
% rng(0);
features = xlsread('features.xlsx');
emPoin   = xlsread('poincare_points.xlsx');
Y = poin2lor(emPoin');

d = 2;
Manifold = hyperbolicfactory(d);
Manifold.transp = @paralleltransp;
Manifold.retr = @Manifold.exp;
Manifold.invretr = @Manifold.log;
Manifold.type = 'HBLM'; % Hyperbolic manifolds represented by Lorenzt models. 

Dw = 4;

numSamples = max(size(Y)) - 1;

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

alpha = zeros(nagents, lncostterms);

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

Z  = {}; % train set of embeddings
WW = {}; % train set of features

for i = 1:numSamples + 1
%     Z{i} = (Y(:,i));
    Z{i} = projection(Y(:,i));
    WW{i} = features(i,:)';
end
indx  = 100;
indxs = [1:indx-1,indx+1:numSamples + 1];
Z  = {Z{indxs}};
WW = {WW{indxs}};
% 'primate' as the test point
w = features(indx,:)';

tau = 1;
tau0 = tau;
tau1 = tau;

epsilon_i = 0.5;
delta_i = 1e-4;
for i = 1:nagents    
    probtmp.M = Manifold;
    probtmp.lncostterms = lncostterms;
    D = {Z{(i-1)*lncostterms+1 : i*lncostterms}};
    W = {WW{(i-1)*lncostterms+1 : i*lncostterms}};
    alpha(i,:) = computAlpha(w, W);
    probtmp = agent_Prob_HSP(D, alpha(i,:), lncostterms, probtmp);
    prob{i} = probtmp;
 
    sigma2(i) = log(1/delta_i) * tau^2 / lncostterms / lncostterms;    
    optstmp.tau  = tau;
    optstmp.tau0 = tau0;
    optstmp.tau1 = tau1;
    opts{i} = optstmp;
%     checkgradient(prob{i});
end


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


options.agent = opts;
% for step size
options.stepsize_type = 'fix';
options.option = 2;
options.st = 1;

options.maxiter = 150;

rng(0)
xopt = projection(Y(:,indx));
Seed = 10;
fvs1   = [];
ngfs1  = [];
infos1 = {};
fvs1_   = [];
ngfs1_  = [];
infos1_ = {};
fvs2   = [];
ngfs2  = [];
infos2 = {};
fvs2_   = [];
ngfs2_  = [];
infos2_ = {};

for i = 1:Seed
    fprintf("\n============ i:%2d ============\n", i)
    rng(i+Seed)
    x0 = Manifold.rand();
    options.localMethod = 'DPRSGD';    
    options.stepsize = .2;    
    % SGD
    fprintf("i:%d, SGD\n", i)
    K = 4;
    for j = 1:nagents
        options.agent{j}.lmaxiter = K;
        options.agent{j}.sigma = 0*sqrt(K * sigma2(j) / epsilon_i / epsilon_i);
        options.agent{j}.batch_size = lncostterms;
        options.agent{j}.stepsize = options.stepsize;
    end
    rng(i+Seed+1)
    [x1, fv1, info1] = PriRFed(problem, x0, options);
    fvs1 = [fvs1, info1.fvs];
    ngfs1 = [ngfs1, info1.num_gfs];
    infos1{i} = info1;    

    % DP-SGD
    fprintf("i:%d, DP-SGD\n", i)
    rng(i+Seed+1)
    for j = 1:nagents
        options.agent{j}.lmaxiter = K;
        options.agent{j}.sigma = sqrt(K * sigma2(j) / epsilon_i / epsilon_i);
        options.agent{j}.batch_size = lncostterms;
        options.agent{j}.stepsize = options.stepsize;
    end
    [x1_, fv1_, info1_] = PriRFed(problem, x0, options);
    fvs1_  = [fvs1_, info1_.fvs];
    ngfs1_ = [ngfs1_, info1_.num_gfs];
    infos1_{i} = info1_;

    % SVRG
    fprintf("i:%d, SVRG\n", i)
    rng(i+Seed+2)
    options.localMethod = 'DPRSVRG';    
    options.stepsize = .2;
    K = 2; m = 3;
    for j = 1:nagents
        options.agent{j}.lmaxiter = K;
        options.agent{j}.sigma = 0*sqrt(m*K * sigma2(j) / epsilon_i / epsilon_i);
        options.agent{j}.batch_size = 10;
        options.agent{j}.lmaxinneriter = m;
        options.agent{j}.stepsize = options.stepsize;
    end
    [x2, fv2, info2] = PriRFed(problem, x0, options);
    fvs2 = [fvs2, info2.fvs];
    ngfs2 = [ngfs2, info2.num_gfs];
    infos2{i} = info2;

    % DP-SVRG
    fprintf("i:%d, DP-SVRG\n", i)
    rng(i+Seed+2)
    K = 2; m = 3;
    for j = 1:nagents
        options.agent{j}.lmaxiter = K;
        options.agent{j}.sigma = sqrt(m*K * sigma2(j) / epsilon_i / epsilon_i);
        options.agent{j}.batch_size = 10;
        options.agent{j}.lmaxinneriter = m;
        options.agent{j}.stepsize = options.stepsize;
    end
    [x2_, fv2_, info2_] = PriRFed(problem, x0, options);
    fvs2_ = [fvs2_, info2_.fvs];
    ngfs2_ = [ngfs2_, info2_.num_gfs];    
    infos2_{i} = info2_;
end


figure(1)
plot(mean(ngfs1, 2), mean(fvs1, 2),'b -', mean(ngfs1_, 2) ,mean(fvs1_, 2),'b -.', mean(ngfs2, 2), mean(fvs2, 2), 'r -', mean(ngfs2_, 2), mean(fvs2_, 2), 'r -.','LineWidth',2)
ylabel('Function value, $f(x^{(t)})$', 'interpreter','latex', 'FontSize', 18)
xlabel('#grad', 'FontSize',18)
legend('PriRFed-SGD-\infty', 'PriRFed-SGD-0.5',...
       'PriRFed-SVRG-\infty', 'PriRFed-SVRG-0.5', 'FontSize',18)


figure(2)
t = [0:1:options.maxiter];
y1 = zeros(Seed, options.maxiter+1);
y2 = zeros(Seed, options.maxiter+1);
y3 = zeros(Seed, options.maxiter+1);
y4 = zeros(Seed, options.maxiter+1);
for i = 1:Seed
    y1(i, 1) = Manifold.dist(xopt,x0);
    y2(i, 1) = Manifold.dist(xopt,x0);
    y3(i, 1) = Manifold.dist(xopt,x0);
    y4(i, 1) = Manifold.dist(xopt,x0);
    for j = 1:length(t)-1
        y1(i,j+1) =  Manifold.dist(xopt, infos1{i}.xs{t(j+1)});
        y2(i,j+1) =  Manifold.dist(xopt, infos1_{i}.xs{t(j+1)});    
        y3(i,j+1) =  Manifold.dist(xopt, infos2{i}.xs{t(j+1)});
        y4(i,j+1) =  Manifold.dist(xopt, infos2_{i}.xs{t(j+1)});    
    end
end
semilogy(mean(ngfs1, 2), mean(y1, 1),'b -', mean(ngfs1_, 2), mean(y2,1), 'b -.', mean(ngfs2, 2), mean(y3,1), 'r -', mean(ngfs2_, 2), mean(y4,1), 'r -.', 'LineWidth',2);
ylabel('Distance to the true point, $\mathrm{dist}(x^{(t)},x_{\mathrm{true}})$', 'interpreter','latex', 'FontSize', 18)
xlabel('#grad', 'FontSize',18)
legend('PriRFed-SGD-\infty', 'PriRFed-SGD-0.5',...
       'PriRFed-SVRG-\infty', 'PriRFed-SVRG-0.5', 'FontSize',18)

[r1, c1] = findMinIndex(y1, 1);
[r2, c2] = findMinIndex(y2, 1);
[r3, c3] = findMinIndex(y3, 1);
[r4, c4] = findMinIndex(y4, 1);

x1  = infos1{r1}.xs{c1-1};
x1_ = infos1_{r2}.xs{c2-1};
x2  = infos2{r3}.xs{c3-1};
x2_ = infos2_{r4}.xs{c4-1};

% plot poincare disk

xopt = lor2poin(xopt);
x1   = lor2poin(x1);
x1_  = lor2poin(x1_);
x2   = lor2poin(x2);
x2_  = lor2poin(x2_);

P = xlsread('poincare_points.xlsx');
np = readcell('names.xlsx');
mc = readcell('edgelist.xlsx');

namep = [];
for i = 1:length(np)
    namep = [namep, string(np{i})];
end

x = [xopt';x1'; x1_'; x2'; x2_'];
namex = [" ", "PriRFed-SGD-\infty", "PriRFed-SGD-0.5", "PriRFed-SVRG-\infty", "PriRFed-SVRG-0.5"];

figure(3)
plot_Poincare_disk(P, x, namep, namex, mc);

hold on
figure(3)
scatter(xopt(1), xopt(2), 100, [102/255 139/255 139/255], 'filled', 'LineWidth',2);
hold on
h1 = scatter(x1(1), x1(2), 150, [139/255 101/255 8/255], 'diamond', 'LineWidth',2);
hold on
h2 = scatter(x1_(1), x1_(2), 150, [139/255 101/255 8/255], 'filled', 'diamond','LineWidth',2);
hold on
h3 = scatter(x2(1), x2(2), 150, [139/255 101/255 8/255], 'square' ,'LineWidth',2);
hold on 
h4 = scatter(x2_(1), x2_(2), 150, [139/255 101/255 8/255], 'filled', 'square', 'LineWidth',2);
% [~, ldg] = legend([h1, h2, h3, h4], "PriRFed-SGD-\infty", "PriRFed-SGD-0.5", "PriRFed-SVRG-\infty", "PriRFed-SVRG-0.5",'FontSize',25);
[~, ldg] = legend([h1, h2, h3, h4], "PriRFed-SGD-\infty", "PriRFed-SGD-0.5", "PriRFed-SVRG-\infty", "PriRFed-SVRG-0.5",'FontSize',25,'Location','eastoutside','Orientation','vertical');
ldg(5).Children.MarkerSize = 12;
ldg(6).Children.MarkerSize = 12;
ldg(7).Children.MarkerSize = 12;
ldg(8).Children.MarkerSize = 12;


% saveas(gca, 'PATHMNISTFM'+string(nagents)+'x'+string(ncostterms))

% figure(1)
% t = [0:options.checkperiod:options.maxiter];
% plot(t ,mean(fvs1'),'-p', t ,mean(fvs2'),'-o', t, fv*ones(length(t), 1), 'k -.','LineWidth',2)
% ylabel('cost', 'FontSize',25)
% xlabel('iteration', 'FontSize',25)
% legend('DP-RSGD', 'DP-RSVRG', 'True Value', 'FontSize',25)
% 
% saveas(gca, 'PATHMNISTFM'+string(nagents)+'x'+string(ncostterms))


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

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

% computing alpha(w)
function output = computAlpha(w, W)   
    N = length(W);
    Kw = zeros(N, 1);
    K  = zeros(N, N);
    for i = 1:N
        Kw(i) = kernel(w, W{i});
        for j = 1:N
            K(i, j) = kernel(W{i}, W{j});
        end
    end
    output = (K + 1e-5 * eye(N)) \ Kw;
end

function output = kernel(w1, w2)
    dn = norm(w1 - w2, 2);
    sigma = 1/(2*0.03^2);
    output = 1 / exp(sigma * dn * dn);
end

% parallel transport
function output = paralleltransp(x, y, u)
    output = u - (inner_product(y, y, u) / inner_product(x, x, y)) * (x + y);
end

function output = projection(u)
    output = u / sqrt(abs(inner_product(u, u, u)));
%     output = u;
end

function output = inner_product(x, v, u)
    output = v' * u - 2 * v(1) * u(1);
end


function [row, col] = findMinIndex(X, dim)
    if nargin == 1
        dim = 1;
    end
    [n, m] = size(X);
    indx = find(X == min(min(X, [], dim)));
    col = ceil(indx / n);
    row = indx - (col - 1) * n;   
end


