function [betas,stats] = weighted_glmfit(Y,varargin)
% Calculate weighted average using weighted linear least squares
% See examples below for usage
%
% Model:
% Y_i = 1*Ypop + noise
%
% :Inputs:
%
%   **Y:**
%        data matrix (nsub x T)
%
%   **w:**
%        weights
%
%   **varY:**
%        variance of data at each time point (nsub x T) + var between
%
% :Outputs:
%
%   **Ymean:**
%        weighted mean of each column of Y
%
%   **dfe:**
%        error degrees of freedom, adjusted for inequality of variance
%        (Sattherwaite) and pooled across data columns
%
% :Extended output in stats structure:
%
%   **stats.t:**
%        t-values for weighted t-test
%
%   **stats.p:**
%        2-tailed p-values for weighted t-test
%
%   **r:**
%        weighted correlation coeff across columns of Y
%
%   **xy:**
%        weighted covariance matrix
%
%   **v:**
%        weighted variance estimates for each column of Y
%          - sqrt(v) is the standard error of the mean (or grp difference)
%
%
%   **stats.fits:**
%        fits  for each group (Ymean by group), low contrast weight group then high
%
% Fastest if no stats are asked for.
%
% :Computation time:
% For FULL stats report
%   - Triples from 500 -> 1000 columns of Y, continues to increase
%
% For mean/dfe only, fast for full dataset (many columns of Y)
%
% :Examples:
%
% Basic multivariate stats for 1000 columns of dat, no weighting
% Multivariate covariances are meaningful if cols of Y are organized, e.g.,
% timeseries
% ::
%
%    [means,stats] = weighted_glmfit(dat(:,1:1000));
%
% The same, but return univariate stats only (good for large Y)
% ::
%
%    [means,stats] = weighted_glmfit(dat,'uni');
%
% A weighted version, where we put in the weights, and with a design matrix too:
% ::
%
%    [means,stats] = weighted_glmfit(X,dat,'uni','w',weights);
%
% A weighted version, where weights are determined from w/i subject variances:
% ::
%
%    [means,stats] = weighted_glmfit(X,dat,'uni','vary',variances);
%
% ..
%    NOTE: TOR CHANGED INPUT TO ASSUME THAT WE SHOULD ENTER VARWI + VARBETWEEN
% ..

if nargin == 0, error('Must at least enter data as 1st argument.'); end

% --------------------------------------
% * Set up arguments
% --------------------------------------

domultivariate = 0;     % multivariate covariance est for Y

zpdiff = []; w = []; varY = []; X = [];

for i = 1:length(varargin)
    arg = varargin{i};
    if ischar(arg)
        switch lower(arg)
            case 'w', w = varargin{i+1};
            case 'vary', varY = varargin{i+1};
            case 'uni', domultivariate = 0;
            case 'multi', domultivariate = 1;
                
            case {'X','x'}, X = varargin{i+1};
        end
    end
end

% fill in missing inputs with default values

if ~is_entered(w), w = ones(m,1);  end
if ~is_entered(varY), varY = ones(m,1); end
if ~is_entered(X), X = ones(m,1);  end

% Get rid of missing values
nancols = find(any(isnan(Y) | Y==0,1));
Y(:,nancols) = [];

[m,n] = size(Y);

% --------------------------------------
% * Means and contrast
% --------------------------------------

% pool weights across all voxels
[betas,invxwx,bform,fits] = get_betas_singleweight(X,Y,w);

if nargout == 1, return, end

% --------------------------------------
% * Residuals
% --------------------------------------

e = Y - fits;         % residuals


% --------------------------------------
% * Degrees of freedom
% --------------------------------------

[dfe,dfediff] = get_dfe(m,n,X,bform,varY,0);



if ~domultivariate
    % ======================================
    %
    %
    % Univariate stats: MSE, t, and p-values
    %
    %
    % ======================================

    % --------------------------------------
    % * Mean squared error
    % --------------------------------------

    W = diag(w);                    % Weight matrix
    % Loop version of MSE: avoids out of memory errors for large voxel sets
    MSE = zeros(1,n); for i=1:n, MSE(i) = e(:,i)'*W*e(:,i); end, MSE = MSE/dfe;

    k = size(betas,1);
    v = repmat(diag(invxwx),1,n) .* repmat(MSE,k,1);
    %v = invxwx * MSE;       % variances for mean

    % output
    stats.descrip1 = 'Univariate stats for test against zero:';
    stats.v = v;
    stats.v_descrip = 'V = ste^2; variance of mean estimate';
    stats.t = betas ./ sqrt(v);
    stats.p = 2 * ( 1 - tcdf(abs(stats.t),dfe) );
    stats.dfe = dfe;



else
    % ======================================
    %
    %
    % Multivariate stats: MSE, cov(Y), r(Y)
    % Useful for simulating t-values under dependence
    %
    % ======================================

    % --------------------------------------
    % * Mean squared error
    % --------------------------------------

    % additional output: covariance matrix for betas and zdiff across time
    % (columns)
    % and correlation matrix for betas and zdiff
    % used in Monte Carlo simulations for controlling false positives
    % across columns

    MSE = (e'*W*e)/dfe;             % Mean square error

    if dobtwn
        MSEdiff = (ediff'*W*ediff)/dfediff;
    end


    % --------------------------------------
    % * Estimated covariance and correlation
    %   Estimated between-subjects variance (v)
    % --------------------------------------

    xy = invxwx * MSE;           % Covariance matrix for betas;

    xy = 0.5*(xy+xy');              % Remove rounding error

    if dobtwn
        xydiff = inv(bcon'*W*bcon)*MSEdiff;           % Covariance matrix for betas;
        xydiff = 0.5*(xydiff+xydiff');
    end

    v = diag(xy);                   % Variance for betas

    if dobtwn
        vdiff = diag(xydiff);
    end

    r = xy./sqrt(v*v');             % Correlation matrix for betas

    if dobtwn
        rdiff = xydiff./sqrt(vdiff*vdiff');         % Correlation matrix for betas
    end

    stats.descrip1 = 'Multivariate stats for test against zero:';
    stats.r = r;
    stats.v = v;
    stats.xy = xy;

    stats.t = betas ./ sqrt(v');
    stats.p = 2 * ( 1 - tcdf(abs(stats.t),dfe) );

end


return








function [dfe,dfediff] = get_dfe(m,n,X,bform,varY,dobtwn,hatdiff,bcon)

dfediff = [];

% Set up residual-forming matrix
% --------------------------------------
dfe_v = zeros(n,1);
R = eye(m) - X*bform;    % residual inducing matrix

% contrast, if entered
if dobtwn
    dfe_vdiff = zeros(n,1);
    Rdiff = eye(m) - bcon * hatdiff;
end

% Calculate effective degrees of freedom
% --------------------------------------

have_unique_vars = size(varY,2) == n;

if ~have_unique_vars

    % Only one (pooled?) vector of variance estimates
    % --------------------------------------
    V = diag(varY(:,1));
    dfe = (trace(R*V)^2)/trace(R*V*R*V);       % Satherwaite approximation
    if dobtwn, dfediff = (trace(Rdiff*V)^2)/trace(Rdiff*V*Rdiff*V); end
else
    % Variance estimates for each data vector
    % --------------------------------------
    for i=1:n,

        % make diagonal matrix of variances
        V = diag(varY(:,i));

        dfe_v(i) = (trace(R*V)^2)/trace(R*V*R*V);       % Satherwaite approximation

        if dobtwn
            dfe_vdiff(i) = (trace(Rdiff*V)^2)/trace(Rdiff*V*Rdiff*V);
        end

    end

    dfe = mean(dfe_v);               % Calculate average df over all columns (pool over data vectors)
    if dobtwn
        dfediff = mean(dfe_vdiff);
    end

end

return





function bool = is_entered(x)

bool = exist('x','var') && ~isempty(x);

return


%
%
% Duplicated in robust_reg_pooled
%
%

function [betas,invxwx,bform,fits] = get_betas_singleweight(X,Y,w)

W = diag(w);                    % Weight matrix

%X = repmat(1,m,1);              % Design matrix - 1 column of all ones to calculate average
% and, separately, use bcon if that's entered

invxwx = inv(X'*W*X);
bform = invxwx * X'* W;         % beta-forming matrix.  hat = X * bform

% rows are columns of design (X), cols are Y variables
betas = bform*Y;

if nargout > 3
    fits = X * betas;
end

return



function w = bisquare_weight(r,radjust,xrank)
% r is residuals
% radjust is adjustment factor: DuMouchel & O'Brien
% xrank is rank of weighted X matrix (design)
% w is weights from bisquare function
% n is number of Y variables to replicate weights over
tuneconst = 4.685;

r = r .* radjust;
s = mad_sigma_pooled(r,xrank);
r = r ./ (s*tuneconst);
w = (abs(r)<1) .* (1 - r.^2).^2;
return


function s = mad_sigma_pooled(r,xrank)
%    Compute std estimate using MAD of residuals from 0
rsort = sort(abs(r));
rsort = rsort(xrank:end,:); % eliminate smallest; like reducing df
s = median(rsort(:)) / 0.6745;
return

function str = display_string(str)
str = sprintf(str); fprintf(1,'%s',str);
return


function erase_string(str)

len = length(str);
str2 = repmat('\b',1,len);

fprintf(1,str2);

return

