cd F:/spcluster/revision/supplementary;
rng(22);

XY = readmatrix("data/syndat.csv");
XX = XY(:, 1:2);
YY = XY(:, 3:3);

N = size(YY, 1);
rng(22);


num_basis = [5^2 10^2 15^2 20^2];
knots_1d = {};
for i = num_basis
    eps = 0; 1/1000 * rand; %0
    %tmp = linspace(0,1,sqrt(i));
    knots_1d{end+1} = linspace(eps,1-eps,sqrt(i));
end;




%Wendland kernel
K = 1; % basis size
Theta = ones(N, sum(num_basis)+1);
for res = 1:length(num_basis)
    theta = 1/sqrt(num_basis(res))*2.5;
    [knots_s1, knots_s2] = meshgrid(knots_1d{res},knots_1d{res});
    knots = [knots_s1(:), knots_s2(:)];
    for i = 1:num_basis(res)
        disp([res i]);
        d = vecnorm(XX -knots(i, :), 2, 2)./theta;
        for j = 1:length(d)
            if d(j) >= 0 & d(j) <= 1
                Theta(j,i + K) = (1-d(j))^6 * (35 * d(j)^2 + 18 * d(j) + 3)/3;
            else
                Theta(j,i + K) = 0;
            end
        end
    end
    K = K + num_basis(res);
end



idx    = find(sum(Theta > 0) >= 50);
Theta  = Theta(:, idx);


x = linspace(0., 1. , 500); 
y = linspace(0., 1. , 500);
[xx,yy] = meshgrid(x,y); 
X = [xx(:) yy(:)];

K = 1; % basis size
Theta2 = ones(size(X,1), sum(num_basis)+1);
for res = 1:length(num_basis)
    theta = 1/sqrt(num_basis(res))*2.5;
    [knots_s1, knots_s2] = meshgrid(knots_1d{res},knots_1d{res});
    knots = [knots_s1(:), knots_s2(:)];
    for i = 1:num_basis(res)
        disp([res i]);
        d = vecnorm(X -knots(i, :), 2, 2)./theta;
        d(d > 1) = 1;
        Theta2(: ,i + K) = (1-d).^6 .* (35 * d.^2 + 18 * d + 3)/3;
    end
    K = K + num_basis(res);
end



Theta2 = Theta2(:, idx);

W = 1/N * Theta2' * Theta2; %'

Q = chol(inv(W));
Q = Q';
Phi = Theta * Q;
W2 = 1/N * Phi' * Phi; %

clear Theta2, X;
max(abs(W2 - eye(size(Theta,2))), [], 'all')








%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  Train-Test %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

rng(22);
sub_idx = 1:N; %randsample(train_idx, 1000, false);

Y = YY(sub_idx, :);
phi = Phi(sub_idx, :);
n = size(Y, 1);




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  MCMC %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


rng(22);

p = size(phi,2);
nu = 2;
A2s = 1*ones(p,1);
%A2s(1) = 5;

niter  = 600;
burnin = 300;
s2s    = zeros(1, niter-burnin);
Omegas = zeros(p, p, niter-burnin);
alphas = zeros(p, niter-burnin);
as     = zeros(p, niter-burnin);
Ys     = zeros(N, niter-burnin);


mu = 0;
alpha = zeros(p,1);
a =  ones(p,1);
tau2 = A2s(1); %10;





for j = 1:niter
    if mod(j, 10) == 0
        disp(j);
    end
    ee = Y - phi * alpha;
    inner_prod = sum((ee - mu).^2);
    s2  = 1/gamrnd(n/2, 2/inner_prod);                     %% s2
    mu  = sum(ee)/(n + s2/tau2) + sqrt(s2/(n+s2/tau2))*randn;                     %% mu
    Omega_term_2 = 2* nu * diag(1./a);
    Omega_scale = inv(Omega_term_2 + alpha*alpha');        %'
    Omega = wishrnd(Omega_scale, nu + p);                  %% Sig_inv
    alpha_Sig = inv(Omega + 1/s2 * (phi' * phi));          %'
    alpha_mu = alpha_Sig * (1/s2) * phi' * (Y - mu);       %'
    alpha = mvnrnd(alpha_mu', alpha_Sig)';                 %%alpha 
    for k = 1:p
        a(k) = 1/gamrnd(        0.5*(nu + p),  1/   (1/A2s(k,1) + nu* Omega(k,k))   );    %% a
    end


    %%%%%%% Store %%%%%%%%%
    

    if j > burnin
        s2s(j-burnin)           = s2;
        Omegas(:,:,j-burnin)    = alpha_Sig; %Omega;
        alphas(:, j-burnin)     = alpha;
        as(:, j-burnin)         = a;
        Ys(:, j-burnin)         = mu + Phi * alpha + sqrt(s2) * randn(N,1);
    end
end


%yhat = Phi * mean(alphas , 2);









