% weighted least squares routine, modified from John Plumb's 
% unweighted least squares code 'fitaline.m'
% inputs:  (X,Y,sigma), aka (index, observation, standard deviation of point)
% outputs:  [m, b, Yans]
%    m - slope of line
%    b - y-intercept
%    Yans - y-values which fit the line calculated, Yans = mX + b

function [m, b, Psq, Yans] = weightline(X,Y,sigma)
%function [m,b,Yans] = weightline(X,Y,sigma)
%function [Yans] = weightline(X,Y)

% get all inputs as column vectors:
if size(X,2)>size(X,1)
  X = X';
end
if size(Y,2)>size(Y,1)
  Y = Y';
end
if size(sigma,2)>size(sigma,1)
  sigma = sigma';
end

m0 = 0;    %initial guess for the slope
b0 = 0;    %initial guess for the y-intercept
%calculate prefit residuals:
m = m0;
b = b0;

A = [ ];  %?x2
% uses least squares to fit a line, so build A-matrix
% with rows [obs-index, 1]:
A(:,1) = X;
A(:,2) = ones(length(X),1);
% state vector (answers):
state = [m;b];   %2x1
% prefit residual = original data:
L = Y;

% build the weight matrix:
W = zeros(length(X), length(X));
for i = 1:length(X)
    W(i,i) = 1/(sigma(i,1)^2);
end

% least-squares solution:
%xhat = inv(A'*A)*A'*L;
xhat = inv(A'*W*A)*A'*W*L;
% Uncertainties:
P = inv(A'*W*A);
Psq = sqrt(P);
% where sig_m = sqrt(P(1,1)) 
% and sig_b = sqrt(P(2,2));
m = xhat(1,1);
b = xhat(2,1);
Yans = m*X + b;