function [ Ppre, params, RMSE, pseed ] = multigaussians( t,P,nk,varargin )
% Companion function to manuscript:
%  "Fitting Gaussian, Hubbert, or other Bell curves to a time series".
%  Submitted to Mathematical Geosciences.
%
% Find bestfitting gaussians or logistic derivatives to data.
%   Default is Gaussian.
%
%   Inputs: t = independant 'time' variable, P is data to fit, and nk is number
%       of Gaussians to fit to the data
%   Outputs: Ppre = model predictions to data (i.e., misfits = |P - Ppre|)
%       params = nkx3 array containing parameters for nk gaussian curves.
%           column 1 = centroids (mu)
%           column 2 = standard deviations (sigma - tau for logistic derivatives)
%           column 3 = amplitudes at centroids (A)
%       i.e., for Gaussian: P = A*exp(-0.5*((t-mu)^2)/sigma^2)
%
%   Can use own time seeds for where peaks lie by entering
%       [ vector of 1 to n elements ] as fourth argument or by 'mu',[ vector ].
%       If other parameters are to also be seeded, enter:
%        'sigma',[ vector ] for centroids and/or 'A',[ vector ] for amplitudes.
%       These will be attached in sequence to time seeds.
%       'random' can be enetered as an argument to generate randomly selected
%       time seeds (other than those explicitly input) - a
%       useful to get away from local minima. Otherwise, time seeds default
%       to either locations in data peaks if 'findpeaks' exists
%       (SigProc toolbox) or to regularly spaced times.
%       'plot' will generate a plot of data and fit.
%
%   Written by James Conder, Aug 15, 2013
%   Adapted for both Gaussians and Logistics Nov, 2013
%   Cleanup for publication, Jan 3, 2013

% for plotting (default is no plotting)
plot_it = false;
tplot = t;      % time vector for plotting purposes
hP = [];

% data & model specs
ndata = length(P);
np = 3;

vectorflipped = false;   % expecting a column vector. flip if row.
if ~iscolumn(P)
    vectorflipped = true;
    P = P';
    t = t';
end
P0 = P;     % save initial data vector in P0

% normalize time to stabilize ai
tstretch = 1./(max(t) - min(t));
t = t*tstretch;
tmin = min(t);
tmax = max(t);
timewidth = tmax-tmin;
dtmin = min(diff(t));

% convergence related variables
maxiter = 1000;	% number of maximum iterations
epsilon = 1;
iter = 0 ;              % initialize counter
tol = sqrt(nk)*1.e-6;    % tolerance for parameter improvement
dmscale = 0.4;      % 0-1. larger = more rapid convergence, smaller = more stable

% use cumulative values as monotonic curves tend to be more stable for non-linear inversion
w = 0*t + tstretch;
w(2:end) = diff(t);
Q = cumsum(w.*P0);
Q0 = Q;

% set up initial values for parameters and other useful constants
swidth = 0.1*timewidth/nk;       % starting std dev widths
ascale = Q(end)/(2.5*nk*swidth);

% look for time seeds from command line
randomseedtimes = false;
gaussian = true;
mi = []; ai = []; si = [];             % intial parameters

if nargin > 3
    if isnumeric(varargin{1})
        mi = varargin{1};
    end
    for i = 1:nargin-3
        if ~isnumeric(varargin{i})
            switch varargin{i}                
                case { 'random', 'Random', 'RANDOM' }
                    randomseedtimes = true;
                    
                case { 'logistic','Logistic','hubbert','Hubbert' }
                    gaussian = false;
                    
                case { 'maxiter' }
                    maxiter = varargin{i+1};    % default 1000
                    
                case { 'mu' }
                    mi = varargin{i+1};
                    
                case { 'A' }
                    ai = varargin{i+1};
                    
                case { 'sigma','tau' }
                    si = varargin{i+1};
                    
                case { 'sscale','swidth' }
                    swidth = swidth*varargin{i+1};
                    ascale = ascale/varargin{i+1};
                    
                case { 'plot' }
                    plot_it = true;                    
            end
        end
    end
end
nmuseeds = min(length(mi),nk);
naseeds = min(length(ai),nmuseeds);
nsseeds = min(length(si),nmuseeds);
mi(nmuseeds+1:end) = [];
ai(naseeds+1:end) = [];
si(nsseeds+1:end) = [];
iseeds = [1:3:naseeds*3 2:3:nmuseeds*3 3:3:nsseeds*3];
%iseeds = [2:3:nmuseeds*3];

if ~isempty(mi)
    mi = mi*tstretch;
    si = si*tstretch;
    
    if nmuseeds > naseeds   % pick starting amplitude to seeded times
      for i = naseeds+1:nmuseeds
          j = find(t < mi(i),1,'last');
          ai(i) = P(j);
      end
    end
end
na = nk-length(ai); ns = nk-length(si);
ra = 1; rs = 1;

i = find(Q < 0.01*Q(end),1,'last');
if isempty(i)
    i = 2;
end
tmind = t(i);   % t at 1% of cumulative. Don't seed prior to this
i = find(Q > 0.99*Q(end),1,'first');
tmaxd = t(i);   % t at 99% of cumulative. Don't seed after this

if randomseedtimes
    tmind = max(tmind,tmin+0.1*timewidth);
    mi(nmuseeds+1:nk) = tmind + (tmaxd-tmind)*rand(nk-nmuseeds,1);
    nr = max(na,ns);
    ra = 0.5 + rand(nr,1);
    rs = 1./ra;
    if na < nr
        ra(na+1:end) = [];
    end
    if ns < nr
        rs(ns+1:end) = [];
    end
else
    if nmuseeds == 0
        if exist('findpeaks','file') == 2   % look for peaks in data
            Psmooth = P0;       % do a little smoothing
            for i = 1:np*(nk^2)
                Psmooth = smoother(Psmooth,0.67);
            end
            [~,locs] = findpeaks(Psmooth,'npeaks',nk,...
                'minpeakdistance',ceil(0.04*length(P0)/nk),'sortstr','descend');
            mi = t(locs);       % assign peak locations
            ai = Psmooth(locs); % use peak height
        end
    end
    if ~isempty(mi)
        tmind = min(tmind,min(mi));
    end
    mi(end+1:nk) = ...
        evenly_distribute([tmind reshape(mi,1,length(mi)) tmin+timewidth],nk-length(mi));
end
if ~isempty(ai)
    if ai(1)*swidth*sqrt(2*pi) > Q(end)
        swidth = swidth*Q(end)/(ai(1)*swidth*sqrt(2*pi));
    end
end
    
ai(end+1:nk) = ascale*ra.*ones(nk-length(ai),1);
si(end+1:nk) = swidth*rs.*ones(nk-length(si),1);
if ~gaussian
    si = 1./si;
    swidth = 1/swidth;
end

if ~iscolumn(mi)
    mi = mi';
end
if ~iscolumn(ai)
    ai = ai';
end
if ~iscolumn(si)
    si = si';
end
mi0 = mi/tstretch; ai0 = ai;
if gaussian
    si0 = si/tstretch; 
else
    si0 = si*tstretch; 
end
mibest = mi; sibest = si; aibest = ai;

if plot_it
    figure(1)
    clf
    plot(tplot,P0,'LineWidth',2)
    hold on
end


% predictions from starting parameters
if gaussian
    [Ppre, Pi] = getPpre_gaussian(t,ai,mi,si);
    Qpre = getQpre_gaussian(t,ai,mi,si);
else
    [Ppre, Pi] = getPpre_logistic(t,ai,mi,si);
    Qpre = getQpre_logistic(t,ai,mi,si);
end

if plot_it
    for i = 1:nk
        plot(tplot,Pi(:,i),'k--')     % plot starting curves
    end
end
Pprebest = Ppre;
SSEbest = sum((P0-Ppre).^2);
pause(0.2)

%%% set up weighting based on equalizing importance of two kinds of data
%%% get data importances
% cdf partial derivatives
if gaussian
    Gcdf = mkG_cumGaussian(t,ai,mi,si);
    Gpdf = mkG_Gaussian(t,ai,mi,si);
else
    Gcdf = mkG_logistic(t,ai,mi,si);
    Gpdf = mkG_logisticderiv(t,ai,mi,si);    
end
    
G = [ Gcdf ; Gpdf];
inan = isnan(G);
G(inan) = 0;
[U,~,~] = svd(G,0);     % Singular Value Decomposition for data importances
dimportances = abs(diag(U*U'));
dimpratio = sum(dimportances(ndata+1:end))/sum(dimportances(1:ndata));  % P/Q
dimpratio0 = dimpratio;

%%% data weights
W = ones(ndata*2,1);
if dimpratio > 1
    W(ndata+1:2*ndata) = 1./dimpratio;
end
W = diag(W);

%%% model weights
Wm = ones(1,nk*3);
i = 1:3:nk*3-2;
Wm(i) = 2;       % slow amplitude changes (eaisest to make) 
%Wm(iseeds) = 2*Wm(iseeds);       % slow amplitude changes (eaisest to make) 

%%% iterate for best fitting parameters
while epsilon > tol
    iter = iter + 1;
    
%%% set up partial derivative matrix
if gaussian
    Gcdf = mkG_cumGaussian(t,ai,mi,si); % cdf partial derivatives
    Gpdf = mkG_Gaussian(t,ai,mi,si);    % pdf partial derivatives
else
    Gcdf = mkG_logistic(t,ai,mi,si); % cdf partial derivatives
    Gpdf = mkG_logisticderiv(t,ai,mi,si);    % pdf partial derivatives
end
G = [ Gcdf ; Gpdf];
G = W*(G.*repmat(Wm,2*ndata,1));    % put model and data weights into G

Wmi = Wm;
wmseeds = false;
if ~isempty(iseeds) && epsilon > 1.e4*tol
    wmseeds = true;
    Wmi(iseeds) = 10*Wmi(iseeds);
end
  
% set up data vector
d(1:ndata,1) = Q0 - Qpre;   % cdf
d(ndata+1:2*ndata,1) = P0 - Ppre; % pdf
d = W*d;

    
%%% solve for parameters
j = find(sum(abs(G),2) < 1.e-8, 1);    % crude, but fast, measure of matrix condition
if isempty(j)
  [ddm,~] = lsqr(G,d);
  dm = dmscale*ddm./Wmi';
else
  [U,S,V] = svd(G,0);				% Singular Value Decomposition
  Sinvdiag = 1./diag(S) ;
  ising = Sinvdiag(1)./Sinvdiag < 1.e-12 ;
  Sinvdiag(ising) = 0;
  Sinv = diag(Sinvdiag);
  dm = (dmscale*V*Sinv*U'*d)./Wmi';

  dimportances = abs(diag(U*U'));
  dimpratio = sum(dimportances(ndata+1:end))/sum(dimportances(1:ndata));  % P/Q
end

normpenalty = 0;
if sum(isnan(dm)) > 0       % abort if NaNs found
    disp('WARNING: badly scaled G-matrix. May have difficulty converging.')
    inan = isnan(dm);
    dm(inan) = 0;
    normpenalty = 1.e6*tol;
end

ai = ai + dm(1:np:np*nk-np+1);
mi = mi + dm(2:np:np*nk-np+2);
si = si + dm(3:np:np*nk-np+3);

% see how model paramaters have changed
  epsilon = norm(dm) + normpenalty;
  if iter == maxiter
	    disp(['max number of iterations (' num2str(maxiter) ') reached...exiting'])
        disp(['normalized epsilon: ' num2str(epsilon/tol)]) 
        epsilon = tol ;
  end
  
  
% predictions from new parameters
if gaussian
    Ppre = getPpre_gaussian(t,ai,mi,si);
    Qpre = getQpre_gaussian(t,ai,mi,si);
else
    Ppre = getPpre_logistic(t,ai,mi,si);
    Qpre = getQpre_logistic(t,ai,mi,si);
end
SSE(iter) = sum((P0-Ppre).^2);

if SSE(iter) < SSEbest
    SSEbest = SSE(iter);
    Pprebest = Ppre;
    mibest = mi; aibest = ai; sibest = si;
end

if sum(isnan(dm)) > 0       % abort if NaNs found
    disp('WARNING: error during convergence. aborting...')
    keyboard
    epsilon = tol;
    SSE(iter) = SSE(1);
end

% plot updated curve
if plot_it
    %figure(1)
    %hold on    
    hP(end+1) = plot(tplot,Ppre,'r--');
    if length(hP) > 20
        delete(hP(1))
        hP(1) = [];
    end
    pause(0.01)

    %dataimportances = diag(U*U');
    %figure(2)
    %clf
    %hold on
    %plot(tplot,dataimportances(1:ndata),'x')
    %plot(tplot,dataimportances(ndata+1:2*ndata),'o')
    %title('data importances')
    %pause(0.01)
    %figure(1)
    %keyboard
    
end

% upweight pdf derivatives relative to cdf during later iterations
if iter < 0.95*maxiter
  W = diag(W);
  W0 = W(end)+0.02;
  W(ndata+1:2*ndata) = W0;
  W = diag(W);
  if epsilon < 1.e3*tol
    W = diag(W);
    W0 = min(W(end)+0.5,max(40,40/dimpratio));
    W(ndata+1:2*ndata) = W0;
    W = diag(W);
    dmscale = min(0.8,1.05*dmscale);
  end
  %if epsilon > 1.e4*tol && W(end) > 1
  %  W = diag(W);
  %  W0 = max(0.5*W(2*ndata),0.5/dimpratio0);
  %  W(ndata+1:2*ndata) = W0;
  %  W = diag(W);
  %  dmscale = max(0.4,0.5*dmscale);
  %end
else
    W = diag(W);
    W0 = W(end)+2;
    W(ndata+1:2*ndata) = W0;
    W = diag(W);
    dmscale = min(0.9,1.05*dmscale);
end    

ibad = find(ai.*si/sum(ai.*si) < 0.01/nk);
if ~gaussian
    ibad = find((ai./si)/sum(ai./si) < 0.01/nk);
end
ibadm = ibad; ibada = ibad; ibads = ibad;
if wmseeds
    ifixed = ibadm <= nmuseeds; ibadm(ifixed) = [];
    ifixed = ibada <= naseeds; ibada(ifixed) = [];
    ifixed = ibads <= nsseeds; ibads(ifixed) = [];
end
mi(ibadm) = 0.5*(tmind+tmaxd) + 0.25*timewidth*randn(length(ibadm),1);
ai(ibada) = 0.75*ascale;
si(ibads) = swidth;

ai = min(ai,2*max(P0));
ibad = find(ai <= 0);
ibadm = ibad; ibada = ibad;
if wmseeds
    ifixed = ibadm <= nmuseeds; ibadm(ifixed) = [];
    ifixed = ibada <= naseeds; ibada(ifixed) = [];
end
ai(ibada) = 0.75*ascale;
mi(ibadm) = 0.5*(tmind+tmaxd) + 0.25*timewidth*randn(length(ibadm),1);

ibad = find(si <= 0);
if ~gaussian
    ibad = find(1./si < 0.5*dtmin);
end
ibadm = ibad; ibads = ibad;
if wmseeds
    ifixed = ibadm <= nmuseeds; ibadm(ifixed) = [];
    ifixed = ibads <= nsseeds; ibads(ifixed) = [];
end
si(ibads) = swidth;
mi(ibadm) = 0.5*(tmind+tmaxd) + 0.25*timewidth*randn(length(ibadm),1);

if gaussian
    ibad = find(si > timewidth);
else
    ibad = find(1./si > timewidth);
end
ibadm = ibad; ibads = ibad;
if wmseeds
    ifixed = ibadm <= nmuseeds; ibadm(ifixed) = [];
    ifixed = ibads <= nsseeds; ibads(ifixed) = [];
end
si(ibads) = swidth;
mi(ibadm) = 0.5*(tmind+tmaxd) + 0.25*timewidth*randn(length(ibadm),1);

ibad = find(mi > tmaxd+0.4*(timewidth));
jneg = find(mi < tmind-0.1*timewidth);
ibad = union(ibad,jneg);
ibadm = ibad;
if wmseeds
    ifixed = ibadm <= nmuseeds; ibadm(ifixed) = [];
end
if ~isempty(ibad)
    mi(ibadm) = 0.5*(tmind+tmaxd) + 0.25*timewidth*randn(length(ibadm),1);
end


end
%disp([num2str(iter) ' iterations. W = ' num2str(W(end))])

% check if final is 'best'
if SSEbest < SSE(end)
    mi = mibest; ai = aibest; si = sibest;
    Ppre = Pprebest;
end
RMSE = sqrt(SSEbest/ndata);

[~,I] = sort(mi);
mi = mi(I); si = si(I); ai = ai(I);
[~,I] = sort(mi0);
pseed = [ mi0(I) si0(I) ai0(I) ];


% output final predicted data and model parameters
if gaussian
    si = si/tstretch;
else
    si = si*tstretch;
end
mi = mi/tstretch;
params = [ mi si ai ];

if vectorflipped
    Ppre = Ppre';
end

if plot_it
    t = t/tstretch;
    figure(1)
    clf
    hold on
    plot(t,P0,'LineWidth',2)
    hold on
    plot(t,Ppre,'r','LineWidth',2)
    
    if gaussian
        [~, Pk]  = getPpre_gaussian(t,ai,mi,si);
        [~, Pk0] = getPpre_gaussian(t,ai0,mi0,si0);
    else
        [~, Pk]  = getPpre_logistic(t,ai,mi,si);
        [~, Pk0] = getPpre_logistic(t,ai0,mi0,si0);
    end

    for k = 1:nk
        plot(t,Pk0(:,k),'c:')
        plot(t,Pk(:,k),'k--')
    end
    box on
end

end

%%%%%%%%%%%%%%%%%% internal functions %%%%%%%%%%%%%%%%%%%%%

function [ G ] = mkG_cumGaussian(t,ai,mi,si)
%Build G-matrix for fit to cumulative Gaussian
%   INPUTS: t = independent (time) variable
%           ai = amplitudes of seed curves
%           mi = time of seed curve peaks
%           si = width (standard deviation) of seed curves

nk = length(mi);
ndata = length(t);
np = 3;
sq2 = sqrt(2);
sqpi = sqrt(pi);

G = zeros(ndata,np*nk);
for i = 1:nk
    z = (t - mi(i))/(sq2*si(i));
    B = ai(i)*sqpi/sq2;
    ddda = si(i)*(sqpi/sq2)*(1 + erf(z));
    dddm = -ai(i)*exp(-z.*z);
    ddds = B - B*erf(z) - (2/sqpi)*B*z.*exp(-z.*z);
    G(:,(i-1)*np+1) = ddda;
    G(:,(i-1)*np+2) = dddm;
    G(:,(i-1)*np+3) = ddds;
end
inan = isnan(G);
G(inan) = 0;
end

function [ G ] = mkG_Gaussian(t,ai,mi,si)
%Build G-matrix for fit to standard Gaussian
%   INPUTS: t = independent (time) variable
%           ai = amplitudes of seed curves
%           mi = time of seed curve peaks
%           si = width (standard deviation) of seed curves

nk = length(mi);
ndata = length(t);
np = 3;

G = zeros(ndata,np*nk);
for i = 1:nk
    z = -0.5*((t - mi(i)).^2)/(si(i)*si(i));
    ddda = exp(z);
    dddm = ai(i)*(t-mi(i)).*exp(z)/(si(i)^2);
    ddds = ai(i)*((t-mi(i)).^2).*exp(z)/(si(i)^3);
    G(:,(i-1)*np+1) = ddda;
    G(:,(i-1)*np+2) = dddm;
    G(:,(i-1)*np+3) = ddds;
end
inan = isnan(G);
G(inan) = 0;
end

function [ G ] = mkG_logisticderiv(t,pmax,mu,tau)
%Build G-matrix for fit to logistic derivative
%   INPUTS: t = independent (time) variable
%           pmax = amplitudes of seed curves
%           mu = time of seed curve peaks
%           tau = width (standard deviation) of seed curves

nk = length(mu);
G = zeros(length(t),3*nk);
for k = 1:nk
    ik = (k-1)*3 ;
    ee = min(exp(-tau(k)*(t-mu(k))),1.e12) ;
    eee = 1./((1 + ee).^2) ;
  
    G(:,ik+1) = 4*ee.*eee;                    % ddp/dpmax
    G(:,ik+2) = 4*pmax(k)*tau(k)*ee.*eee + ...
        -8*pmax(k)*tau(k)*ee.*ee.*eee./(1 + ee);                 % ddp/dmu
    G(:,ik+3) = -4*pmax(k)*(t-mu(k)).*ee.*eee + ...
       + 8*pmax(k)*(t-mu(k)).*ee.*ee.*eee./(1 + ee);       % ddp/dtau
end
inan = isnan(G);
G(inan) = 0;
end


function [ G ] = mkG_logistic(t,pmax,mu,tau)
%Build G-matrix for fit to logistic
%   INPUTS: t = independent (time) variable
%           pmax = amplitudes of seed curves
%           mu = time of seed curve peaks
%           tau = time constant of curves

nk = length(mu);
Qinf = 4*pmax./tau;
G = zeros(length(t),3*nk);
for k = 1:nk
    ik = (k-1)*3 ;
    ee = min(exp(-tau(k)*(t-mu(k))),1.e12) ;
    eee = 1./((1 + ee).^2) ;
    G(:,ik+1) = 4./(tau(k)*(1 + ee)) ;                % ddq/dpmax
    G(:,ik+2) = -Qinf(k)*tau(k)*(ee.*eee) ;         % ddq/dmu
    G(:,ik+3) = Qinf(k)*(t-mu(k)).*ee.*eee ;    % ddq/dtau
end
inan = isnan(G);
G(inan) = 0;
end


function [Ppre, Pk] = getPpre_gaussian(t,ai,mi,si)
nk = length(mi);
Pk = zeros(length(t),nk);
for k = 1:nk
    Pk(:,k) = ai(k).*exp(-0.5*((t-mi(k)).^2)/(si(k)*si(k)));
end
Ppre = sum(Pk,2);
end

function [Qpre, Qk] = getQpre_gaussian(t,ai,mi,si)
nk = length(mi);
Qk = zeros(length(t),nk);
sq2 = sqrt(2);
sqpi = sqrt(pi);

for k = 1:nk    
    z = (t - mi(k))/(sq2*si(k));
    Qk(:,k) = ai(k)*si(k)*(sqpi/sq2)*(1 + erf(z));
end
Qpre = sum(Qk,2);
end


function [Ppre, Pk] = getPpre_logistic(t,pmax,mu,tau)
nk = length(mu);
Pk = zeros(length(t),nk);
for k = 1:nk  
    ee = min(exp(-tau(k)*(t-mu(k))),1.e12);
    eee = 1./((1 + ee).^2) ;
    Pk(:,k) = 4.*pmax(k)*ee.*eee;  % predicted P
end
Ppre = sum(Pk,2);
end

function [Qpre, Qk] = getQpre_logistic(t,pmax,mu,tau)
nk = length(mu);
Qinf = 4*pmax./tau;
Qk = zeros(length(t),nk);
for k = 1:nk  
    ee = min(exp(-tau(k)*(t-mu(k))),1.e12);
    Qk(:,k) = Qinf(k)./(1 + ee) ;  % predicted Q
end
Qpre = sum(Qk,2);
end