         % multivariate case sv + dpm sampler.
% prior is Normal Wishart, but can be used for any exponential conjugate
% families.
% code for the cluster parameters update is only for univariate case.
clear; clc;
%{
load 'USCPI.csv';
y = USCPI;
load 'hhat_SVM.mat';
hhat_SVM = hhat;
load 'hCI_SVM.mat';
hCI_SVM = hCI;
%}

dataraw = importdata('crspvwr.csv');
data1 = dataraw.data(1:6815);
data2 = dataraw.data(1:end);
data3 = dataraw.data(6816:end);
y = 100*data1;

T = length(y);

burnin = 120;
nloop = 320;
%% prior
dlt0 = .97; Vdlt = .1^2;
nusig2 = 10; Ssig2 = .2^2*(nusig2-1);
a0 = 2; b0 = 8;
%a0 = .01; b0 = .01;
mu0 = 0; kappa0 = 10; nu0 = 10; S0 = 10;
prior = normalWishartSV(kappa0,mu0,nu0,S0);
model_name = 'DPM+SV';
%warning('off')
%% initialize the Markov chain

dlt = .98; signu2 = .2^2;
h = sqrt(signu2)*randn(1,T);
hh = h'; exph = exp(h);
alpha = 1;
X = y'; [d,~] = size(X);
[z,Theta,w] = oneDPM(X,h,alpha,prior); nk = T*w; llh = zeros(nloop,1);

%% initialize for storeage
store_h = zeros(nloop - burnin,T);
store_lam2 = zeros(nloop - burnin,T);
store_eta = zeros(nloop - burnin,T);
store_z = zeros(nloop - burnin,T);
store_phi = zeros(nloop - burnin,5);

store_clstpara = [];

%% compute a few things outside the loop
Hdlt = speye(T) - sparse(2:T,1:(T-1),dlt*ones(1,T-1),T,T);
newnusig2 = T/2 + nusig2;

counth = 0; countac0 = 0; countac1 = 0;
disp(['Starting MCMC for ' model_name '.... ']);
start_time = clock;

for loop = 1:nloop
   
    %% sample assignment parameters z
    % update new h for Theta
    %{
    if length(nk) ~= length(Theta)
        error('nk and Theta are not equal in length')
    end
    %}
    for j = 1:length(Theta)
        idxT = find(z==j);
        Theta{j} = prior.addData(X(:,idxT),h(idxT));
    end
   
    for i = randperm(T)
        x = X(:,i);
        hi = h(i);
        k = z(i);
        Theta{k} = Theta{k}.delSample(x,hi);
        nk(k) = nk(k)-1;

        Pk = log(nk)+cellfun(@(t) t.logPredPdf(x,hi), Theta);
        P0 = log(alpha)+prior.logPredPdf(x,hi);
        p = [Pk,P0];
        llh(loop) = llh(loop)+sum(p-log(T));
        p = p - max(p,[],2);
        p = exp(p);
        p = p./sum(p,2);
    
        k = randmn(p);
        if k == numel(Theta)+1                 % add extra cluster
            Theta{k} = prior.clone.addSample(x,hi);
            nk = [nk,1];
        else
            Theta{k} = Theta{k}.addSample(x,hi);
            nk(k) = nk(k)+1;
        end
        z(i) = k;   % update z
        
        idx0 = find(nk==0); % remove empty cluster
        %{
        if length(idx0) > 1
            msg = 'idx0 has multiple values';
            error(msg)
        end
        %}
        if ~isempty(idx0)
            Theta(idx0) = [];
            nk(idx0) = [];
            which = z>idx0;
            z(which) = z(which)-1;
        end     
    end
    
    K = length(Theta);
    ot = sum(nk < T*0.01); % number of small clusters
    
    lam2K = []; etaK = [];
    for i = 1:K
        Ui = Theta{i}.U_;
        kappai = Theta{i}.kappa_;
        mui = Theta{i}.mu_;
        nui = Theta{i}.nu_;
        S = Ui'*Ui - kappai*(mui*mui');
        cholS = chol(S);
        invcholS = cholS\speye(d);
        lam2K(i) = wishrnd(invcholS*invcholS',nui,invcholS);
        etaK(i) = mui + chol(kappai*lam2K(i))'\randn(d,1);
    end
    
    zK = sparse([1:T]',z',ones(T,1),T,K);   % update univariate case cluster parameters
    lam2 = zK*lam2K'; eta = zK*etaK';         
    etalam = eta.*sqrt(lam2);
    alp = zeros(T,1);
    
    %% sample h
    yl = y.*sqrt(lam2); 
    HinvSH = Hdlt'*sparse(1:T,1:T,[(1-dlt^2)/signu2; 1/signu2*ones(T-1,1)])*Hdlt;
    deltah = zeros(T,1);
    HinvSHdeltah = HinvSH*deltah;
    s2 = (yl-etalam).^2;
    errh = 1; ht = hh;
    while errh> 10^(-3)
        expht = exp(ht);
        sinvexpht = s2./expht;
        alp2expht = alp.^2.*expht;
        fh = -.5 + .5*sinvexpht - .5*alp2expht;
        Gh = .5*sinvexpht + .5*alp2expht;
        Kh = HinvSH + spdiags(Gh,0,T,T);
        newht = Kh\(fh+Gh.*ht+HinvSHdeltah);
        errh = max(abs(newht-ht));
        ht = newht;          
    end 
    cholHh = chol(Kh,'lower');
    % AR-step:     
    hstar = ht;
    uh = hstar-deltah;
    logc = -.5*uh'*HinvSH*uh -.5*sum(hstar) + ...
       - .5*exp(-hstar)'*(yl-etalam-alp.*exp(hstar)).^2 + log(3);
  
    flag = 0;
    while flag == 0
        hc = ht + cholHh'\randn(T,1);
        vhc = hc-ht;
        uhc = hc-deltah;
        alpARc = -.5*uhc'*HinvSH*uhc -.5*sum(hc) + ...
            -.5*exp(-hc)'*(yl-etalam-alp.*exp(hc)).^2 + .5*vhc'*Kh*vhc - logc; 
        countac0 = countac0 + 1;
        if alpARc > log(rand)
            flag = 1;
            countac1 = countac1 + 1;
        end
    end        
    % MH-step
    vh = hh-ht;
    uh = hh-deltah;
    alpAR = -.5*uh'*HinvSH*uh -.5*sum(hh) + ...
        -.5*exp(-hh)'*(yl-etalam-alp.*exp(hh)).^2 + .5*vh'*Kh*vh - logc;
    if alpAR < 0
        alpMH = 1;
    elseif alpARc < 0
        alpMH = - alpAR;
    else
        alpMH = alpARc - alpAR;
    end    
    if alpMH > log(rand) || loop == 1
        hh = hc;
        exph = exp(hh);
        counth = counth + 1;
    end
    h = hh';    % update h
    
    %% sample dlt signu2
    % sample signu2
    errh = [hh(1)*sqrt(1-dlt^2); hh(2:end)-dlt*hh(1:end-1)];    
    newSsig2 = Ssig2 + sum(errh.^2)/2;    
    signu2 = 1/gamrnd(newnusig2, 1./newSsig2); 
    % sample dlt
    Vdlthat = 1/(1/Vdlt + hh(1:end-1)'*hh(1:end-1)/signu2);
    dlthat = Vdlthat*(dlt0/Vdlt + hh(1:end-1)'*hh(2:end)/signu2);
    dlt = tnormrnd(dlthat,Vdlthat,-1,1);
    
    %% sample alpha

    xi = beta(alpha+1,T);
    pi = (a0+K-1)/((a0+K-1)+T*(b0-log(xi)));
    if rand < pi
        alpha = gamrnd(a0+K, 1/(b0-log(xi)));
    else
        alpha = gamrnd(a0+K-1, 1/(b0-log(xi)));
    end
  
    if loop>burnin
        i = loop-burnin;
        store_h(i,:) = h';
        store_z(i,:) = z;
        store_lam2(i,:) = lam2';
        store_eta(i,:) = eta';
        store_phi(i,:) =  [alpha K dlt signu2 ot]; 
        clstpara = [etaK; lam2K; nk];
        store_clstpara = [store_clstpara,clstpara];
    end    
    if (mod(loop,50) == 0)
        disp([num2str(loop) ' loops... '])
    end 
end
clc;
disp( ['MCMC takes '  num2str( etime( clock, start_time) ) ' seconds' ] );

phihat = mean(store_phi)'
phistd = std(store_phi)'
hhat = mean(store_h)';
etahat = mean(store_eta)';
lam2hat = mean(store_lam2)';

hCI = quantile(store_h,[.05 .95])';
lam2CI = quantile(store_lam2,[.05 .95])';
phiCI = quantile(store_phi,[.05 .95])';

%{
tid = linspace(1948,2013.25,T)';    % compare the results to Chan 17
figure;    
hold on    
    plot(tid, hhat, 'LineWidth',2,'Color','black');  
    plot(tid, hCI,'--r','LineWidth',2);  
    plot(tid, hhat_SVM, 'LineWidth',2,'Color','blue');
    plot(tid, hCI_SVM,'--c','LineWidth',2);
    legend('hhat','hCI','hCI','hhatSVM','hCISVM','hCISVM');
hold off    
xlim([min(tid)-1 2014]);  title('h_t');  
%}    

tid = linspace(2007,2021,T)';
figure;    
hold on    
    plot(tid, hhat, 'LineWidth',2,'Color','black');  
    plot(tid, hCI,'--r','LineWidth',2);  
hold off    
xlim([min(tid)-1 2022]);  title('h_t'); 

% y = 100*data1;
%{
tid = linspace(1980,2007,T)';
figure;    
hold on    
    plot(tid, hhat, 'LineWidth',2,'Color','black');  
    plot(tid, hCI,'--r','LineWidth',2);  
hold off    
xlim([min(tid)-1 2008]);  title('h_t'); 
%}