% MAR_Bayes.m, built by B.X. Semmens based on modified code from
%MDSsearchAIC.m written by M.D. Scheuerell. Last update 6/23/08

% requires MLR_Gibbs_func.m

% This program does MCMC fitting of multivariate AR(1) using Gibbs sampling

% ->  Based on code originally written by AR Ives

%******************************************
% initialization routines
clear all

% Global variables
global P Q R ib jb mb nb
global lagstate statevar covariate 
global conlimit bootn
global rawvar rawcovar
global varind covarind

%******************************************


%******************************************
% LOAD YOUR DATASET

load lwa62to94.txt /ascii;
data = lwa62to94;
%******************************************


%******************************************
% ASSIGN COLUMN NAMES
% you can use anything here; I picked these for clarity
% if you change a name, just search & replace it throughout the program
year = data(:,1);       % dummy variable coding for blocks of data to be considered continuous in the obs error subprogram; often years for limnological data, but could be experimental treatments, for example
time = data(:,2);       % time step number 
driver1 = data(:,3);    % month
driver2 = data(:,4);    % month squared a la tony ives to control for season
driver3 = data(:,5);    % temp 20 to 0 m avg
driver4 = data(:,6);    % TP, lots interpolated -- see notes 12, 13 july
driver5 = data(:,7);    % pH
compt1 = data(:,8);     % Cryptomonas bv
compt2 = data(:,9);     % Diatom edible bv
compt3 = data(:,10);    % Green edible bv
compt4 = data(:,11);    % bluegreen bv
compt5 = data(:,12);    % Unicells bv
compt6 = data(:,13);    % Other algae bv
compt7 = data(:,14);    % Conochilus #
compt8 = data(:,15);    % cyclops #
compt9 = data(:,16);    % daphnia #
compt10 = data(:,17);   % diaptomus #
compt11 = data(:,18);   % epischura #
compt12 = data(:,19);   % leptodora #
driver6 = data(:,20);   % Neomysis #/L
compt13 = data(:,21);   % non-daphnid cladocerans (bosmina diaphanosoma) #
compt14 = data(:,22);   % non-colonial rotifers #/l

% Set the number of compartments & number of drivers in your dataset
ncompts = 14;
ndrivers = 6;

%******************************************
% MCMC model function parmeters
            iters = 10000;
            burn = 1000;
            thin = 10;
% MCMC chain storage matrices
        E = zeros(iters/thin,ncompts); %posterior error chain
        A = zeros(iters/thin,ncompts); %posterior slope chain
        B = zeros(ncompts,ncompts,iters/thin); %post. interact coeff
        C = zeros(ncompts,ndrivers,iters/thin); %post. covar coeff

%******************************************
% DEFINE WHICH VARIATES TO INCLUDE
rawvar = [compt1 compt2 compt3 compt4 compt5 compt6 compt7...
          compt8 compt9 compt10 compt11 compt12 compt13 compt14];

% DEFINE WHICH COVARIATES TO INCLUDE
rawcovar = [driver1 driver2 driver3 driver4 driver5 driver6];
%******************************************


%******************************************
% PREPARE DATA: 

% Get sizes of matrices
[nr nc] = size(rawvar);
[nrc ncc] = size(rawcovar);

% If desired, transform rawvar and/or rawcovar, e.g. by a log-transform
% For Steph's data set in this analysis -- data are log-transformed before going in
% rawvar=log(rawvar);	
	
% LAG THE DATA INTO SEPARATE MATRICES AT TIME T AND TIME T-1
s1 = length(rawvar(:,1));
XX = [rawvar(1:s1-1,:) rawvar(2:s1,:)];

% GET COVARIATES
% Do you want to lag the covariates? If so, set lag=1. Otherwise, set lag=0.
lag = 1;
CC = rawcovar(2-lag:s1-lag,:);
    
% THROW OUT NON-OVERLAPPING YEARS
XX = XX(find(year(1:s1-1,1)==year(2:s1,1)),:);
lagstate = XX(:,1:nc);
statevar = XX(:,nc+1:2.*nc);
covariate = CC(find(year(1:s1-1,1)==year(2:s1,1)),:);
%******************************************

	
%*****************************************************************
% THESE NEXT TWO LINES SHOULD BE TAILORED TO YOUR DATASET:
% set up index matrices for the variates and covariates 
% zero elements => not included in model
% one elements  => included in model


% set up an index matrix for the interactions among variates
% e.g., indexBGlobal=[1 0.5 0.5; 0.5 1 0.5; 0.5 0.5 1];

% for the L WA dataset...
%               cr  di  gr  bg  uni oth con cyc dph dip epi lep nd  nc
indexBGlobal = [1.0 1 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 1.0 1.0;...   % crypt
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 1.0 1.0;...   % diatom
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 1.0 1.0;...   % green
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 1.0 1.0;...   % blue-green
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 1.0 1.0;...   % unicell
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 1.0 1.0;...   % other algae
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % cono
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % cyclops
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % daph
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % diapt
                0.0 0.0 0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % epi
                0.0 0.0 0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % lepto
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0;...   % non-daphnid clad
                1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0];     % non-cono rotifers

% indexBGlobal = ones(ncompts); %use this to include every interaction

% set up an index matrix for the covariate effects on each variate
% e.g., indexCGlobal=[1 0; 0 1; 0 1];

% for the L WA dataset...
%               mo  mo2 tmp tp  pH  neomysis
indexCGlobal = [1.0 1.0 1.0 1.0 1.0 0.0;... % crypt
                1.0 1.0 1.0 1.0 1.0 0.0;... % diatom
                1.0 1.0 1.0 1.0 1.0 0.0;... % green
                1.0 1.0 1.0 1.0 1.0 0.0;... % bg
                1.0 1.0 1.0 1.0 1.0 0.0;... % unicell
                1.0 1.0 1.0 1.0 1.0 0.0;... % other algae
                1.0 1.0 1.0 0.0 1.0 1.0;... % cono
                1.0 1.0 1.0 0.0 1.0 1.0;... % cyclops
                1.0 1.0 1.0 0.0 1.0 1.0;... % daphnia
                1.0 1.0 1.0 0.0 1.0 1.0;... % diaptomus
                1.0 1.0 1.0 0.0 1.0 1.0;... % epischura
                1.0 1.0 1.0 0.0 1.0 1.0;... % lepto
                1.0 1.0 1.0 0.0 1.0 1.0;... % nondaphnid
                1.0 1.0 1.0 0.0 1.0 1.0];   % non-cono rotifer

% indexCGlobal = ones(ndrivers); %use this to include every interaction

% initialize variates
Q = length(statevar(:,1));      % length of the time series
P = length(statevar(1,:));      % number of variates in model
R = length(covariate(1,:));     % number of covariates in model

hitB = zeros(P,P);
hitC = zeros(P,R);

bestGlobalB = zeros(P,P);
bestGlobalC = zeros(P,R);

%*****************************************************************

% MAIN BODY

for dv=1:P, % loop over the # of variates (e.g. spps) in the model
       
        varind=indexBGlobal(dv,:);  % index of variates to include
        covarind=indexCGlobal(dv,:);% index of covariates to include
        e = zeros(Q,1);             % residuals
        Yhat = zeros(Q,1);          % estimates
                
        Y = statevar(:,dv);         % dependent variates
        lv = find(varind(1,:)==1);  % cell refs to variates to use
        cv = find(covarind(1,:)==1);% cell refs to covariates to use
        nvar = sum(varind(1,:));
        X = [ones(Q,1) lagstate(:,lv) covariate(:,cv)];	% predictors (includes intercept)
                
        % least squares parameter estimates (inits for Gibbs)
        beta = (X'*X)\(X'*Y);
        % calculate residuals for that dependent variate
        e(:,1) = Y - X*beta;
        sigma = e'*e/Q; % estimate of the variance
        
        % BEGIN GIBBS STEP HERE *******************************
            %set up priors
            % prior marix --
            %1) first row is prior on variance (gam(A,B))
            %2) subsequent rows are priors on mean (first col) and
            %   precision (second col)
            priors = zeros(2+sum(varind)+sum(covarind),2); % all prior means are 0 (uninform case)
            priors(1,:) = .01; %prior for model precision (gam(A,B))
            priors(2:2+sum(varind)+sum(covarind),2) = 1000; %prior variance for the B  and C values (uninform case)
            
            %set up inits
            inits = [1/sigma; beta];

            %now do the MCMC
            tic
            posterior=MLR_Gibbs_func(Y,X,inits,priors,burn,thin,iters);
            toc
        % END GIBBS STEP HERE **********************************
        
        % Output MCMC parameter chains...
        E(:,dv) = posterior(:,1);
        A(:,dv) = posterior(:,2);
        for i = 1:sum(varind)
        B(lv(i),dv,:) = posterior(:,2+i); 
        end
        for i= 1:sum(covarind)
        C(cv(i),dv,:) = posterior(:,sum(varind)+2+i) ; 
        end
  
end; % main loop over # of variates