function [cum_p,output,misclass] = dnn(training_data,test_data,beta,rho,loop,burn,loop_aux,u,type)

% INPUTS:
% training_data and test_data.
% k = number of classes.
% beta is the interaction param; rho is the distance parameter; 
% loop = overall no of iterations; burn = burn-in no of iterations;
% loop_aux = no of auxiliary chain iterations;  u=1 corresponds to
% pseudolikelihood estimation, otherwise exchange algorithm.
% type takes cases "norm", "cut-off" or "exp" corresponding to 
% dnn_1, dnn_2 or dnn_3, respectively. (See accompanying paper for 
% more details.

% OUTPUTS:
% cum_p = estimated probability of class label for each feature.
% output = output from Markov chain for beta and rho parameters
% misclass = misclassification rate.

dim_y = size(test_data,1); 
dim_x = size(training_data,1); 

dim_y2 = size(test_data,2);
dim_x2 = size(training_data,2);

x = training_data(:,1:dim_x2-1); % training features
group = training_data(:,dim_x2); % training class labels
y = test_data(:,1:dim_y2-1); % test features
test = test_data(:,dim_y2); % test class labels

k = max(test);

[neigh_class,neigh_dist,rho_max] = nearest_neighs_k(x,group);
[neigh_classy,neigh_disty] = nearest_neighsy_k(y,x,group);

cum_p = zeros(dim_y,k);

output = zeros(loop,2);

[wt] = dist2wt(neigh_dist,rho,dim_x,type);

accept=0;

for l=1:loop
    
    [wty] = dist2wt_y(neigh_disty,rho,dim_x,type);
    
    [p] = predict_class_k(k,neigh_classy,wty,beta,dim_y);
    
    if( l>burn ) 
        cum_p=cum_p + p; 
    end
    
    if(u==1)
        [beta,rho,accept] = update_pseudo(k,beta,rho,rho_max,group,...
            neigh_class,neigh_dist,neigh_class1,dim_x,accept);
    else
        [beta,rho,accept,wt] = exchange_move(k,beta,rho,neigh_dist,...
            group,neigh_class,dim_x,loop_aux,accept,wt,type);
    end
    
    output(l,1) = beta;
    output(l,2) = rho;
    if( mod(l,500)==0 )
        save -ascii cum_p.dat cum_p
        save -ascii output.dat output
        disp(['loop= ' num2str(l) ' beta= ' num2str(beta) ' rho= ' num2str(rho) ' accept= ' num2str(accept/l) ]);
    end
        
end

cum_p = cum_p/(loop-burn);

[c] = cum_p2map_p(cum_p,k);
misclass = 1 - sum(test==c)/dim_y;
disp(['dnn misclassification rate = ' num2str(misclass) ]);

accept_rate=accept/loop;
disp(['acceptance rate = ' num2str(accept_rate) ]);

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

function [p] = predict_class_k(k,neigh_classy,wty,beta,dim_y)

p = zeros(dim_y,k);

for i=1:dim_y

    predict = zeros(1,k);

    for j=1:k
        
        tmp = neigh_classy(i,:)==j;
        predict(j) = exp(beta*sum(wty(i,:).*tmp));
        
    end
    
    predict = predict/sum(predict);

    p(i,:) = predict;

end

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

function [beta,rho,accept] = update_pseudo(k,beta,rho,rho_max,group,neigh_class,neigh_dist,neigh_class1,dim_x,accept)

beta2 = normrnd(beta,0.2); 
[rho2,r1] = unif_prop(rho,0.05,rho_max,0.1);

prior_beta = normpdf(beta,2,2);
prior_beta1 = normpdf(beta2,2,2);

neigh_dist1 = neigh_dist<rho2;
neigh_class2 = neigh_dist1.*neigh_class;

[prob] = pseudolikelihood(k,beta,group,neigh_class1,dim_x);
[prob_prop] = pseudolikelihood(k,beta2,group,neigh_class2,dim_x);

alpha = prob_prop + log(prior_beta1) + log(r1) - prob - log(prior_beta);
if( alpha>=log(rand) )
    beta = beta2;
    rho=rho2;
    accept=accept+1;
end

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

function [prob] = pseudolikelihood(k,beta,group,neigh_class1,dim_x)

prob=0;

for i=1:dim_x
    tot = zeros(k,1);
    tot_neighs = sum(neigh_class1(i,:)~=0);
    if( tot_neighs==0 )
        prob1 = log(1/k);
    else
        for j=1:k
            tot(j) = sum(neigh_class1(i,:)==j);
        end
        g = group(i);
        prob1 = log( exp(beta*tot(g)/tot_neighs)/(sum( exp(beta*tot/tot_neighs) )) );
    end
    prob = prob + prob1;
end

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

function [beta,rho,accept,wt] = exchange_move(k,beta,rho,neigh_dist,...
    group,neigh_class,dim_x,loop_aux,accept,wt,type)

beta2 = normrnd(beta,0.4); 

[rho2,r1] = unif_prop(rho,0.001,100,0.5);

[wt2] = dist2wt(neigh_dist,rho2,dim_x,type);

[group2,neigh_class2] = gibbs_draw_k1(k,dim_x,beta2,loop_aux,wt2,group,neigh_class);

[prob] = piofx(neigh_class,group,beta,dim_x,wt);
[prob1] = piofx(neigh_class,group,beta2,dim_x,wt2);
[prob_prop] = piofx(neigh_class2,group2,beta2,dim_x,wt2);
[prob_prop1] = piofx(neigh_class2,group2,beta,dim_x,wt);

prior_beta = normpdf(beta,2,40);
prior_beta1 = normpdf(beta2,2,40);

alpha = prob1+prob_prop1 - prob-prob_prop + prior_beta1-prior_beta + log(r1);

if( alpha >= log(rand) )
    beta = beta2;
    rho = rho2;
    accept=accept+1;
    wt=wt2;
end

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

function [wt2] = dist2wt(neigh_dist,rho2,dim_x,type)

switch(type)
    
    case {'norm'}
        wt2 = normpdf(neigh_dist,0,rho2);
        row_tot = sum(wt2,2);
        tmp = kron(row_tot,ones(1,dim_x));
        wt2 = wt2./tmp;
    case {'exp'}
        wt2 = exppdf(neigh_dist,rho2);
        row_tot = sum(wt2,2);
        tmp = kron(row_tot,ones(1,dim_x));
        wt2 = wt2./tmp;    
    case {'cut-off'}
        wt2 = 0.999999*(neigh_dist<rho2) + 0.0000001;%*(neigh_dist<200);
        row_tot = sum(wt2,2);
        tmp = kron(row_tot,ones(1,dim_x));
        wt2 = wt2./tmp;
    case {'tri-cut-off'}
        neigh_dist1 = neigh_dist<rho2;
        neigh_dist2 = neigh_dist1.*neigh_dist; % thresholded distances
        wt2 = (rho2 - neigh_dist2).*neigh_dist1;
        wt2 = 0.999999*(wt2) + 0.0000001;%*(neigh_dist<200);
        
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
function [wt2] = dist2wt_y(neigh_disty,rho2,dim_x,type)

switch(type)
    
    case {'norm'}
        wt2 = normpdf(neigh_disty,0,rho2);
        row_tot = sum(wt2,2);
        tmp = kron(row_tot,ones(1,dim_x));
        wt2 = wt2./tmp;
    case {'exp'}
        wt2 = exppdf(neigh_disty,rho2);
        row_tot = sum(wt2,2);
        tmp = kron(row_tot,ones(1,dim_x));
        wt2 = wt2./tmp;
    case {'cut-off'}
        wt2 = 0.999999*(neigh_disty<rho2) + 0.000001;%*(neigh_disty<200);
        row_tot = sum(wt2,2);
        tmp = kron(row_tot,ones(1,dim_x));
        wt2 = wt2./tmp;
    case {'tri-cut-off'}
        neigh_dist1 = neigh_disty<rho2;
        neigh_dist2 = neigh_dist1.*neigh_disty;
        wt2 = (rho2 - neigh_dist2).*neigh_dist1;
        wt2 = 0.999999*(wt2) + 0.000001;%*(neigh_disty<200);

end

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

function [prob] = piofx(neigh_class,group,beta,dim_x,wt)

tmp = kron(group,ones(1,dim_x));
tmp1 = (neigh_class==tmp);
tot = sum(sum(wt.*tmp1));
prob = beta*tot;

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

function [neigh_class,neigh_dist,rho_max] = nearest_neighs_k(x,group)

% x = training data; k = # of classes; rho = cut-off dist
% neighs = matrix where column i corresponds to the neighbours of test point
% i among the training set.

dim_x = size(x,1);

big_dist=100000;

neigh_dist = zeros(dim_x,dim_x);

for i=1:dim_x    
    
    x_tmp = kron( x(i,:),ones(dim_x,1));
    dist_i = sum( (x_tmp-x).^2,2 ).^(0.5);
    neigh_dist(i,:) = dist_i';
    neigh_dist(i,i) = big_dist;

end

neigh_class = kron(group',ones(dim_x,1));
tmp = ones(dim_x,dim_x)-eye(dim_x,dim_x);
neigh_class = neigh_class.*tmp;
                                           % these lines find a reasonable maximal distance.
neigh_dist1 = neigh_dist - big_dist*eye(dim_x);             % a reasonable maximal radius is the median distance
distances = reshape(neigh_dist1,dim_x^2,1);  % between all possible pairs of points.
%s = find(distances); % Pick out non-zero distances.
%rho_max = median(s(distances));
dis = sort(distances);
rho_max = dis(round(size(dis,1)*0.75)); % 75% percentile of distances

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

function [neigh_class,neigh_dist] = nearest_neighsy_k(y,x,group)

% x = training data; k = # of classes; rho = distance parameter
% neighs = matrix where column i corresponds to the neighbours of test point
% i among the training set.

dim_x = size(x,1); 
dim_y = size(y,1); 

neigh_dist = zeros(dim_y,dim_x); 
neigh_class = zeros(dim_y,dim_x); 

for i=1:dim_y
    for j=1:dim_x

        p = euclid_dist(y(i,:),x(j,:));
        neigh_dist(i,j) = p;
        neigh_class(i,j) = group(j);

    end
end

function [p] = euclid_dist(x,y)

p = sum((x-y).^2).^(0.5);

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

function [group1,neigh_class] = gibbs_draw_k1(k,dim_x,beta,loop,wt,group,neigh_class)

% x = locations of the points. 
% beta = interaction parameter
% rho = cut-off distance
% loop = # of mcmc iterations

group1=group;

for i=1:loop
    for j=1:dim_x

        group_tmp = ceil(k*rand);
        while( group_tmp == group1(j))
            group_tmp = ceil(k*rand);
        end

        tot_prop = wt(j,:).*(neigh_class(j,:)==group_tmp);
        tot_cur = wt(j,:).*(neigh_class(j,:)==group1(j));

        alpha = beta*sum(tot_prop-tot_cur);

        if(log(rand) <= alpha)
            group1(j) = group_tmp;
            neigh_class(:,j) = group1(j)*ones(dim_x,1);
            neigh_class(j,j)=0;
        end


    end
end

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

function [c] = cum_p2map_p(cum_p,k)

size_p = size(cum_p,1);
c1 = zeros(size_p,k);
for i=1:size_p
    c1(i,:) = cum_p(i,:)==max(cum_p(i,:));
    if(sum(c1(i,:))==k)
        c1(i,:)=zeros(1,k);
        cen=ceil(k*rand);
        c1(i,cen) = 1;
    end
end

c_tmp = kron(1:k,ones(size_p,1));
c_tmp = c_tmp.*c1;

c = sum(c_tmp,2);