function [ePlus, niters, exitflag] = calcPositiveEmissions(y, N, xGuess, ...
                                       thresh, param, maxIters, verbose, ...
                                       OptTol, lsqrIters, prob2, ...
                                       eps, maxNonlinIters) 
%calcPositiveEmissions Function to impose non-negativity on the emissions
%   [ePlus, iters, exitflag] = calcPositiveEmissions(y, N, thresh, param, ...
%                                                 maxIters, versbose, OptTol, ...
%                                                 lsqrIters, prob2)
% is a subroutine that imposes non-negativity on fluxes. It assumes that
% the guessed emissions (which may have +ve and -ve components) is pretty
% close the the real emissions i.e., x0 = abs(xGuess), which would be
% non-negative, is a close over-estimate of the emissions. Thereafter, we
% seek ePlus = diag(Corr) * x0, where Corr(i) = exp(xi(i))is a positive correction
% around 1.0 i.e. xi(i) ~ 0. We seek these xi(i). We also stipulate that
% the vast bulk of x0 will not be changed i.e. xi(i) are mostly zeros
%
% Let ePlus_m be the mth guess of ePlus. ePlus_0 = x0. Then we seek xi such
% that |y - H * ePlus_m| is minimized i.e. min( | y - H * ePlus_m | ) 
%
% y - H * ePlus_{m} = y - H * ( diag( exp( xi ) ) * ePlus_{m-1} )
%                   = y - H * ( diag( 1  + xi ) * ePlus_{m-1}
%                   =  y - H * ePlus_{m-1} - H * z,     where z = xi *_e_* ePlus_{m-1}, 
%                                                             *_e_* means element-wise multiplication
%                   =  r - H * z,                       r = y - H * ePlus_{m-1}
%
% So, minimize( |y - H * ePlus_m| ) by changing as few emissions as
% possible (since # of corrections > number of observations) becomes the
% same as as StOMP solution of r - Hz
%
% Thereafter, in an iterative loop
%   xi = z /_e_/ ePlus_{m-1} where /_e_/ is element-wise division
%   xi = sign(xi) * min(abs(xi), 0.5) ; The linearization really works for xi ~ 0
%   Corr = exp(xi)
%   ePlus_m = diag(Corr) * ePlus_{m-1}
%
% Stopping criterion for iterations
%  - Calculate residual | y - H * ePlus_m | / |y|
%  - Calculate change in emissions | ePlus_m - ePlus_{m-1} | / |ePlus_m|
%  - Check if both are small
%
% INPUTS :
%   y      : vector of Co2 concentration observations at towers
%   N      : Dimensionality of the problem; how many emissions to estimate
%   xGuess : a guess at what the emissions may be. Should be close
%   thresh, param, maxIters, verbose, OptTol, lsqrIters : parameters for
%            StOMP
%   prob2  : struct with info on the problem
%   eps    : ratio of change in x to current x that indicated convergence
%   maxNonlinIters : # of iters to go for the linearized solve
%
% OUTPUTS :
%  ePlus : non-negative emissions
%  iters : # of loops taken to impose non-negativity
%  exitflag : if I fail to exist in 100 iters, -1, else 0
%
% Jaideep Ray, 8/23/2012
% -------------------------------------------------------------------------

%% ---- Create a guess for the emissions and calculate residuals
x0 = abs(xGuess(:)) ;
small = 1.0e-12 ;
converged = false ;
niters = 1 ;

%% Iteration loop
while (converged == false) 
    
    % ---- Calculate the residual
    small_resid = false ;
    resid = y - getPred(prob2.H, x0) ;  
    
    if ( norm(resid) / norm(y) < eps )
        small_resid = true ;
    end
    
    % ---- Solve for corrections
    % Use the standard version of StOMP. For the optimized version, contact
    % Jaideep Ray, jairay@sandia.gov
    % [z, iters] = SolveStOMP_JRmod2(@projectionsUsingHmatrices, resid, ...
    %                                N, thresh, param, maxIters, ...
    %                                verbose, OptTol, lsqrIters, prob2) ;
    [z, iters] = SolveStOMP(@projectionsUsingHmatrices, resid, N, thresh, ...
                            param, maxIters, verbose, OptTol) ;
     xi = z ./ (x0 + small) ;
     xi = sign(xi) .* min(abs(xi),  0.5) ; % elements of xi should not exceed 0.5
     corr = exp(xi) ; % multiplicative corrections; should vary between 0.6 and 1.6

     % ---- Updated emissions; are we close yet?
     xnew = x0 .* corr ;
     
     small_deltax = false ;
     if ( norm(xnew - x0) / norm(xnew) < eps )
         small_deltax = true ;
     end

     % ---- All done
     fprintf('calcPositiveEmissions(): iter = %d, norm(resid) = %e, norm(dx) = %e\n', ...
             niters, norm(resid), norm(xnew - x0)) ;
     fprintf('calcPositiveEmissions(): iter = %d, ratio(resid) = %e, ratio(dx) = %e\n', ...
             niters, norm(resid)/norm(y), norm(xnew - x0)/norm(xnew)) ;
     niters = niters + 1 ;

     % ---- Convergence criterion
     if ( (small_resid == true) || (small_deltax == true) )
         converged = true ;
         exitflag  = 0 ;
         ePlus = xnew ;
     end
         
     % ---- Check for non-convergence
     if (niters == maxNonlinIters) % Time to bail
         exitflag = -1 ;
         ePlus = xnew ;
         fprintf('calcPositiveEmissions(): non-convergence!! Reached maxIters\n') ;
         return ;
     end
     
     % ---- If we are here, time to start a new iteration
     x0 = xnew ;
end

end

% =========================================================================

function y = getPred(Hmatrices, x)
% getPred(H, x) returns CO2 concentrations at towers given H matrices and
% emisisons
%  y = getPred(H, x) 
% 
% INPUTS :
%  H : cell array, nTowers long, of H matrices for each tower
%  x : a vector of emissions, should be size(H{i}, 2) long
%  y_j = H{j} * x, where y_j are CO2 concentrations at tower j, H is the
%  corresponding H matrix.
%
% OUTPUTS :
%  y : a vector of co2 concentrations at all towers, at 3 hr resolution
%
% Jaideep Ray, 8/23/2012
% -------------------------------------------------------------------------

nTowers = length(Hmatrices) ;
nObs  = size(Hmatrices{1}, 1) ;
y = zeros( nTowers*nObs, 1) ;

for iTower = 1 : nTowers
    st = 1 + (iTower - 1) * nObs ;
    en = iTower * nObs ;
    y(st:en) = Hmatrices{iTower} * x ;
end

end
% =========================================================================

function [ y ] = projectionsUsingHmatrices(mode, m, n, x, I, dim, UserData)
%projectionUsingHmatrices Computes y = A(:,I)*x if mode = 1, or y = A(:,I)'*x if
%mode = 2
%   [ y ] = projectionUsingHmatrices(mode, m, n, x, I, dim, UserData) computes
%   the projections y = A(:,I)*x or y = A(:,I)'*x. Here x are the emissions
%   (for entire period) and I denotes the emissions (grid-cells) that have been
%   selected by the StOMP in this round of the iteration. 
%   
%   INPUTS :
%   mode : 1 or 2, depending on whether we'll use the H or H' matrix
%   m    : m by dim is the implicit A matrix. Also, length(y) = m in mode 1
%   n    : The number of emissions selected in this iter
%   I    : Indices of the emissions selected. length(I) = n
%   x    : Selected emissions, length(x) = n, if mode 1. Observation
%          residuals, length(x) = m, if mode = 2.
%   dim  : dimensionality of the problem. Total number of emissions, which 
%          need to be sparsified.
%   UserData : Struct, made by the user, and passed into this function
%              untouched by StOMP.
%  
%   OUTPUT :
%   y    : if mode = 1, co2 concetration predictions corresponding to the
%          guessed, sparse emissions sent in. length(y) = m
%          if mode = 2, projections of the residuals onto the H matrix. length(y) = n
%
% -------------------------------------------------------------------------

%% See is there is a UserData. If not, use global version
global PROB_GLOBAL ;
if (nargin < 7)
    UserData = PROB_GLOBAL ;
end

%% Checks

% max value of I should be <= ngrid-cells * nPeriods
assert( max(I(:)) <= size(UserData.H{1}, 2) ) ;

% dimensionality should be equal to number of emissions
assert( dim == size(UserData.H{1}, 2) ) ;

%% Work

if ( mode == 1 ) % Given wavelet coeffs, return co2 concentrations

    % ---- m should be equal to the number of observations
    assert(m == ...
        UserData.info.nObsPerPeriod*UserData.info.nPeriods*UserData.info.nTowers) ;

    % ---- number of emisisons/corrections should be less than those we started with
    assert( length(x) <= (size(UserData.H{1}, 2)*UserData.info.nPeriods) ) ;

    % ---- Allocate output
    y = zeros(m, 1) ; % concetrations over nTowers and nPeriod

    % ---- Loop over towers
    for itower = 1 : UserData.info.nTowers
        H = UserData.H{itower} ;
        Hpart = H(:, I) ;
        tmp = Hpart * x ; % concentrations at this tower due to indicated emission for nPeriods
        len = length(tmp) ;
        iend = itower * len ;
        y(iend - len + 1 : iend) = tmp ;
    end
    return ;
end

if ( mode == 2 ) % Given residuals at each tower, return wavelet coeffs
    y = zeros(n, 1) ; % These will contain wavelet coeffs
    
    % ---- number of residuals == number of observations 
    assert(length(x) <= ....
        (UserData.info.nObsPerPeriod*UserData.info.nPeriods*UserData.info.nTowers)) ;

    nobsPerTower = UserData.info.nObsPerPeriod * UserData.info.nPeriods ;
    for i = 1 : UserData.info.nTowers
        towerObs = x(i*nobsPerTower - nobsPerTower + 1 : i*nobsPerTower) ;
        H = UserData.H{i} ;
        Hpart = H(:, I) ;
        tmp = transpose(Hpart) * towerObs ; % tmp contains contribs to wavelet coeffs from each tower
        y = y + tmp ;
    end
    return ;
end

end