function [beta, t, pvals, convals, con_t, con_pvals, sigma, Phi, df, stebeta, conste, F] = fit_gls(y, X, c, p, varargin)
% Fit a linear model using generalized least squares and an AR(p) model
%
% :Usage:
% ::
%
%     [beta, t, pvals, convals, con_t, con_pvals, sigma, Phi, df, stebeta, conste, F] = fit_gls(y,X,c,p,[PX, equal to pinv(X), for speed], [Weights])
%
% This program uses the Cochrane-Orcutt algorithm to iteratively find the
% GLS solution and estimate the noise parameters.
%
% Step 1: Find the OLS solution.
%
% Step 2: Use residuals from the previous fit to estimate the parameters in
% the AR(p) noise model.
%
% Step 3: Find the GLS solution using the covariance matrix corresponding
% to an AR(p) model with parameters estimated in Step 2 inserted.
%
% Step 4: Repeat steps 2-3 until convergence.
%
% :Inputs:
%
%   **y:**
%        fMRI time course (T x 1 vector)
%
%   **X:**
%        Design matrix (T x param matrix)
%
%   **c:**
%        contrast vector(s) (param x # contrasts matrix)
%
%   **p:**
%        order of AR model.
%
%   **PX:**
%        pinv(X), for speeded, repeated calculations with different y vectors
%
%        Note: if using weights, px = inv(X' * W * X) * X' * W; 
%              where W = diag(Weights);
%
%   **Weights:**
%        Optional, vector of weights for each observation
%
%        Empty or missing: Unweighted analysis.
%
% Note that setting p=0 implies a white noise model.
%
% :Output:
%
%   **t:**
%        t-value for the contrast c'beta
%
%   **df:**
%        degrees of freedom using Satterthwaite approximation
%
%   **beta:**
%        beta vector
%
%   **Phi:**
%        vector of coefficients in AR(p) model
%
%   **sigma:**
%        standard deviation
%
%   **stebeta:**
%        standard error of betas
%
% ..
%    by Martin Lindquist
%
%    Last updated: 3/29/08, Tor Wager, added weighted least squares
%                        verified that beta and t-values are identical to
%                        glmfit.m in matlab7.5 with ar p = 0
%                        ***AR(p) with weighted least squares needs to be
%                        checked.  behaving reasonably.
%               4/1/08,  Tor : output stats for boht betas and contrasts
%                        Reorganized order of outputs
%               5/15/08  Tor : Weird things happening with single inputs; particulaly with aryule; force
%                         double
% ..

y = double(y);
X = double(X);

T = length(y);      % Length of time course

k = size(X, 2);     % predictors


if length(varargin) > 1
    w = double(varargin{2});        % weights for weighted least squares
else
    w = ones(T, 1);
end
W = diag(w);
sqrtW = sqrt(W);     % for weighted residuals

if ~isempty(varargin) && ~isempty(varargin{1})
    px = double(varargin{1});
else
    invxvx = inv(X' * W * X);   % we can re-use this later if p == 0
    px = invxvx * X' * W; 
end


% Step 1: Find the OLS solution 
beta = px*y;                                      % Beta values; weighted, if weights are used
resid = sqrtW * y - sqrtW * X * beta;             % Residuals (Weighted, if weights are used)
sigma = sqrt((1 / (T - k)) * resid' * resid);     %sum(resid.^2)));  % Estimate of Sigma

% Weighted residuals: Three equivalent ways
% We pick one that is compatible with AR estimation
% 1)
% beta = px*y;                       % Beta values
% resid = y - X*beta;                     % Residuals
% sigma = sqrt((1 / (T - k)) * resid' * W * resid); %sum(resid.^2)));  % Estimate of Sigma

% 2) 
%r2 = sqrt(W) * resid; sigma2 = sqrt((1 / (T - k)) * r2' *  r2)

% 3)
% r2 = sqrt(W) * y - sqrt(W) * X * beta;
% sigma2 = sqrt((1 / (T - k)) * r2' *  r2)

% Stuff needed for future iterations
iV = W;  % for ar p = 0 case
A = W;   % for ar p = 0 case
Phi = 0;
betaold = 0;

% Steps 2-4: Find the GLS solution by iteration 
% If p=0, skip this step. Appropriate solution already calculated above.
% Continue iteration until convergence or for at most 10 loops.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Note to Jack and Tor: Keith Worsley uses a similar algorithm when fitting
% a GLM with an AR(p) noise model. However, he skips the iterative step
% and only goes through the loop one time. He claims that this is enough. I
% am not entirely convinced, therefore it is probably better to go through
% a few times if needed.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

i=1;                % Set counter

while i < 10 && p > 0 && (i == 1 || sum((beta - betaold).^2) > 0.001)
    % Do up to 10 iterations, if arp > 0 and either first iteration or
    % there's a difference from last iteration
    
    % resid = y - X*beta;             % Calculate residuals of current model fit

    resid = sqrtW * y - sqrtW * X * beta;             % Residuals (Weighted, if weights are used)

    % Estimate AR parameters using residuals
    [a,e] = aryule(resid, p);
    Phi =zeros(length(a)-1,1);
    Phi(1:p) = -a(2:(p+1));

    %sigma = sqrt(e); % swrtW*e??  % Moved later, because never used until
    %after iteration loop

    % Find the inverse of the covariance matrix
    A = sqrtW; % ***should be sqrt(W)?

    for j=1:p
        %A = A + diag(-Phi(j)*ones((T-j),1),-j);

        A = A + sqrtW * diag(-Phi(j)*ones((T-j),1),-j);
    end;

    %create_figure('A'); imagesc(A, [-.2 .2]); colorbar, drawnow, input(' ')

    iV = A*A';                      %  New weights, with AR estimates, The inverse of the covariance matrix


    betaold = beta;                 % Set old solution to be betaold
    beta = inv(X'*iV*X)*X'*iV*y;    % Calculate new solution

    i = i+1;                        % Add one to counter

    %create_figure('COV'); imagesc(iV, [-.02 .02]); colorbar, drawnow, input(' ')
end

if p > 0 
    % re-calc sigma
    sigma = sqrt(e);

    invxvx = inv(X'*iV*X);
    
    % Should we use the kind of thing below? Seems like sqrt(e) is unweighted,
    % though it seems reasonable...
    % sqrtiv = sqrt(iV);
    % resid = sqrtiv * y - sqrtiv * X * beta;             % Residuals (Weighted, if weights are used)
    % sigma = sqrt((1 / (T - k)) * resid' * resid)
end



R = (eye(T) - X * invxvx * X' * iV);        % Residual inducing matrix

Wd = R * A * inv(iV) * A';                  % inv(iV) = Covariance matrix
df = (trace(Wd).^2)./trace(Wd*Wd);       % Satterthwaite approximation for degrees of freedom


% Should check on below: this creates difference in t-values from
% glmfit...but why wouldn't we need to re-calculate a (larger) sigma if we have reduced df?
% if df ~= (T - k)                    % Tor added 3/29, have to re-calculate sigma
%     sigma = sqrt((1 / df) * resid' * iV * resid); %sum(resid.^2)));  % Estimate of Sigma
% end

varbeta = sigma^2.* invxvx;            % Var(beta)
stebeta = diag(varbeta).^.5;
t = beta ./ stebeta;

convals = [];
conste = [];
con_t = [];
con_pvals = [];
F = [];
if ~isempty(c)
    % Contrast(s)
    convals = (c' * beta);
    conste = diag(c' * varbeta * c).^.5;                   % tor added as output
    con_t = convals ./ conste;                            %sqrt(c'*varbeta*c);             % t-value
end

% get rid of nuisance regressors that aren't in any contrast
% wh = stebeta > 0;
% c = c(wh,wh); 
% beta = beta(wh);
% stebeta = stebeta(wh);
% 
% if ~isempty(c)
%     beta = (c' * beta);
%     %beta = beta(wh);
%     t = beta ./ stebeta;   %sqrt(c'*varbeta*c);            % t-value
% else
%     t = beta(wh) ./ stebeta;
% end

%%%%%% Test H0: beta_1 = beta_2 = .... = beta_param = 0
if nargout > 6
    SSE = y'*y - beta'*X'*y;                 % Error sum of squares
    mSSE = SSE/df;
    J = ones(T);
    SST=y'*y - (1/T).*y'*J*y;           % Total sum of squares
    SSM=SST-SSE;                        % Model sum of squares

    dfSSM=length(c) - 1;                % degrees of freedom for model (param - 1)
    mSSM=SSM/dfSSM;

    F=mSSM/mSSE;         % F-statistic - compare with F-distruibution with (param-1, df) degrees of freedom
end

% % get contrast values if we need those
% if ~isempty(c)
%     beta = (beta' * c)';
% end

if nargout > 2
    pvals = 2 .* (1 - tcdf(abs(t), df)); % two-tailed
    
    
    % make sure p-values for valid results are not zero...
    pvals(pvals == 0 & beta ~= 0 & ~isnan(beta)) = 1000*eps;
    
    if ~isempty(c) && nargout > 5
        con_pvals = 2 .* (1 - tcdf(abs(con_t), df)); % two-tailed
        
        con_pvals(con_pvals == 0 & convals ~= 0 & ~isnan(convals)) = 1000*eps;
    end
        
end

return

