 function [C,GOF] = IRLS_fit(Y, X, ww, typeF)
% Iteratively Reweighted Least Squares (IRLS) with optional fixed point weights (such as water volumes, etc.).
% IRLS iteratively assigns residual weights that are adjusted to downweight data with unusually large residuals.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% Version 1.3  build 2020.06.21
% Author: James Kirchner and Julia Knapp, ETH Zurich
% This code was translated from R to MATLAB by Julia Knapp and tested for
% MATLAB R2018b and R2019b
%
% Copyright (C) 2020 ETH Zurich and James Kirchner
% Public use of this script is permitted under GNU General Public License 3 (GPL3); for details see <https://www.gnu.org/licenses/>
% 
%
% READ THIS CAREFULLY:
% ETH Zurich, James Kirchner and Julia Knapp make ABSOLUTELY NO WARRANTIES OF ANY KIND, including NO WARRANTIES, expressed or implied, that this software is
%    free of errors or is suitable for any particular purpose.  Users are solely responsible for determining the suitability and
%    reliability of this software for their own purposes.
%
% ALSO READ THIS:
% These scripts implement the ensemble hydrograph separation approach as presented in J.W. Kirchner, "Quantifying new water fractions
%    and transit time distributions using ensemble hydrograph separation: theory and benchmark tests", Hydrology and Earth System Sciences, 
%    23, 303-349, 2019.  These scripts are further described in J.W. Kirchner and J.L.A. Knapp, "Matlab and R scripts for ensemble hydrograph
%    separation", Hydrology and Earth System Sciences.  
%    Users should cite both of these papers (the first for the method, and the second for the scripts)


% Y is a numeric vector representing a response variable
% X is a numeric vector or array representing one or more explanatory variables
% ww is an optional numeric vector of point weights (such as masses for mass-weighted regressions)
%    the default weight vector gives all ww's the value of 1
% typeF is an optional string constant indicating the weight function to use.  If typeF is spedified, it must be 'bisquare', 'Welsch', or 'Cauchy'. Anything else generates an error.
% the default weight function is Cauchy

% Y, ww and X (or each vector comprising X, if X is two-dimensional) must be of the same length, otherwise an error results.
% Y and X must not be exactly collinear (that is, there must be some nonzero residuals).  This is not checked.

% IRLS returns a list with two objects:
% - fit: object of class 'lm', generated by a call lm(Y ~ X, weights=wt, na.action='na.omit') 
%        with the iteratively determined weights. This object can then be handled just like any other return from lm, but needs to be prefaced with .fit.
%        Also, the variable names in this 'lm' object will be Y and X (as defined internally here) regardless of what names were passed to this function!
% - wt:   vector of IRLS point weights.

if nargin < 3 % ww not specified
    ww = ones(size(Y));
end
if nargin < 4   % typeF not specified -> use 'Cauchy' (default)
    typeF = 'Cauchy';
end


if ~any(strcmp(typeF,{'bisquare','Welsch','Cauchy'}))
    error('IRLS stopped: no valid weight type specified.  Valid functions are "bisquare", "Welsch", and "Cauchy"')
end
 
if (length(Y) ~= length(ww)) 
	error('IRLS stopped: supplied weight vector must be same length as Y')
end
  
% wwwt = ww.*wt;
wt = ones(size(Y)); % initialize residual weights
wt_chg = 999.0; %initialize weight change
iter = 0; %initialize the iteration counter
rsq = 0;    % initialize r-squared


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%Here's the iteration loop, which runs until the largest weight change for any point is less than 0.01, or iteration limit is exceeded
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
while ((max(wt_chg) > 0.01) && ~all([(iter>10), rsq>0.999]) )
    if (iter>2000) 
        error('IRLS stopped: more than 2000 iterations, sorry!')
    end
    % This error can arise when Y is perfectly collinear with X for more than half the points,
    % and thus the MAR fluctuates near zero, with the weights never stabilizing.
    % That should normally be handled by the r-squared criterion for loop exiting as defined above.
    
    iter = iter+1; %increment the iteration counter
    old_wt = wt; %save the old vector of weights for comparison with the next one
    
    
    if iter == 1
        % start with a line through the median of the data
        ymed = nanmedian(Y);
        resid = Y-ymed;
    else
        
        resid = OUT.residuals;
 
% original code:        const = fit.Coefficients.Estimate(1);   %this is the intercept
% original code:        slope = fit.Coefficients.Estimate(2:length(fit.Coefficients.Estimate)); %this is the vector of regression coefficients
%     
% original code:        %explicit calculation of residuals
% original code:        if (length(slope)==1) 
% original code:            resid = Y - const - X*slope; 
% original code:        else
% original code:            resid = (Y - const - X * slope);
% original code:        end
       
%original code:         can't use residuals(fit) because missing values will mess up the assignment of weights in the steps that follow
%original code:         note %*% is matrix multiplication in R

    end
    
    abs_resid = abs(resid);     
    MAR = nanmedian(abs_resid(abs_resid>0));% remove all zeros
    
    if (MAR==0.0) 
        error('IRLS stopped. Solution has collapsed: median absolute residual is zero!')
    end
    
    if strcmp(typeF,'bisquare') 
        wt = bisquare(abs_resid,MAR); %use bisquare weights
    elseif strcmp(typeF,'Welsch') 
        wt = Welsch(abs_resid,MAR); %use Welsch weights
    elseif strcmp(typeF,'Cauchy') 
        wt = Cauchy(abs_resid,MAR); %use Cauchy weights
    end
    
   
    
    wwwt = ww.*wt;
    
%     A = [X,wwwt];
%     fit = fitlm(A,Y,'y ~ x1');
%     original code: fit = fitlm(X,Y, 'weights',wwwt, 'RobustOpts','off'); %initial least-squares fit
   
   fo = fitoptions('Method','NonlinearLeastSquares',...
                   'Lower',[0,0,-Inf],...
                   'Upper',[Inf,2*pi,0],...
                   'MaxFunEvals',1000,...
                   'MaxIter',1000,...
                   'Startpoint',[1 1 1]);
        
   ft = fittype('A*sin(2*pi*x - Fi)+ K','independent', 'x', 'dependent', 'Y','options',fo);
   
   [C,GOF,OUT] = fit(X,Y,ft,'Weights',wwwt);
  
   % original code: rsq = fit.Rsquared.Ordinary;
   rsq = GOF.rsquare;
   
   wt_chg = abs(wt-old_wt); %calculate change in weights from previous iteration
end 
end




function [w] = bisquare(x, MAR) 
%define bisquare weight function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if (MAR==0) 
    w = zeros(size(x));
    w(x==0) = 1;
else 
  w = (1-(x./(6*MAR)).^2).^2;     %bisquare function
  w(x>6*MAR) = 0;               %replace with zero whenever X is more than 6 times the median absolute residual
end
end


function [w] = Welsch(x, MAR) 
%define Welsch weight function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if (MAR==0) 
    w = zeros(size(x));
    w(x==0) = 1;
else
    w = exp(-(x./(4.4255*MAR)).^2);  %Welsch weight function
end
end


function [w] = Cauchy(x,MAR) 
%define Cauchy weight function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if (MAR==0) 
    w = zeros(size(x));
    w(x==0) = 1;
else
    w = 1./(1+(x./(3.536*MAR)).^2);  %Cauchy weight function
end
end



