function Mh = hmxHorzcat(varargin)
%+========================================================================+
%|                                                                        |
%|         OPENHMX - LIBRARY FOR H-MATRIX COMPRESSION AND ALGEBRA         |
%|           openHmx is part of the GYPSILAB toolbox for Matlab           |
%|                                                                        |
%| COPYRIGHT : Matthieu Aussal (c) 2017-2018.                             |
%| PROPERTY  : Centre de Mathematiques Appliquees, Ecole polytechnique,   |
%| route de Saclay, 91128 Palaiseau, France. All rights reserved.         |
%| LICENCE   : This program is free software, distributed in the hope that|
%| it will be useful, but WITHOUT ANY WARRANTY. Natively, you can use,    |
%| redistribute and/or modify it under the terms of the GNU General Public|
%| License, as published by the Free Software Foundation (version 3 or    |
%| later,  http://www.gnu.org/licenses). For private use, dual licencing  |
%| is available, please contact us to activate a "pay for remove" option. |
%| CONTACT   : matthieu.aussal@polytechnique.edu                          |
%| WEBSITE   : www.cmap.polytechnique.fr/~aussal/gypsilab                 |
%|                                                                        |
%| Please acknowledge the gypsilab toolbox in programs or publications in |
%| which you use it.                                                      |
%|________________________________________________________________________|
%|   '&`   |                                                              |
%|    #    |   FILE       : hmxHorzcat.m                                  |
%|    #    |   VERSION    : 0.42                                          |
%|   _#_   |   AUTHOR(S)  : Matthieu Aussal                               |
%|  ( # )  |   CREATION   : 14.03.2017                                    |
%|  / 0 \  |   LAST MODIF : 15.07.2018                                    |
%| ( === ) |   SYNOPSIS   : H-Matrix horizontal concatenation             |
%|  `---'  |                                                              |
%+========================================================================+

%%% Input analysis
Ml = varargin{1};
Mr = varargin{2};
if (nargin == 2)
    if isequal(Ml.pos{1},Mr.pos{1})
        Ix = (1:size(Ml,1))';
        Iy = (1:size(Ml,2)+size(Mr,2))';
    else
        error('hmxHorzcat.m : unavailable case');
    end
else
    Ml = varargin{1};
    Mr = varargin{2};
    Ix = varargin{3};
    Iy = varargin{4};
end

%%% Preparation
% Particles
X = Ml.pos{1}(Ix,:);
Y = [Ml.pos{2} ; Mr.pos{2}];
Y = Y(Iy,:);

% Accuracy
tol = max(Ml.tol,Mr.tol);

% Initialiation
Mh = hmx(X,Y,tol);

% Admissibility
[isfar,Xdim,Ydim] = hmxFar(Mh);

% Y Indices
bool = (Iy<=size(Ml,2));
Iyl  = Iy(bool);
Iyr  = Iy(~bool)-size(Ml,2);

% Compression for far distances
if isfar
    % Compression
    [Al,Bl] = lowrank(Ml,Ix,Iyl);
    [Ar,Br] = lowrank(Mr,Ix,Iyr);

    % Concatenation
    A = [Al,Ar];
    B = zeros(size(A,2),length(Iy));
    B(1:size(Al,2),bool)      = Bl;
    B(size(Al,2)+1:end,~bool) = Br;

    % Recompression
    [A,B] = hmxQRSVD(A,B,Mh.tol);

    % Validation
    flag = (numel(A)+numel(B)) <= prod(size(Mh));

else
    flag = 0;
end

%%% Compression
if flag
    % Type
    Mh.typ = 1;

    % Low-rank
    Mh.dat = {A,B};

%%%% Full or sparse for smallest box (stopping criterion)
elseif (min(size(Mh)) < 100)
    % Type
    Mh.typ = 2;

    % Full
    Mh.dat = zeros(size(Mh));

    % Fill data
    Mh.dat(:,bool)  = full(Ml,Ix,Iyl);
    Mh.dat(:,~bool) = full(Mr,Ix,Iyr);

%%% H-Matrix (recursion)
else
    % Type
    Mh.typ = 0;

    % Subdivision for X
    [I1,I2] = hmxSubdivide(X,Xdim);
    Mh.row  = {I1 , I1 , I2 , I2};

    % Subdivision for Y
    [I1,I2] = hmxSubdivide(Y,Ydim);
    Mh.col  = {I1 , I2 , I1 , I2};

    % H-Matrix (recursion)
    for i = 1:4
        % Coordinates
        Ir = Mh.row{i};
        Ic = Mh.col{i};

        % Recursion
        Mh.chd{i} = hmxHorzcat(Ml,Mr,Ix(Ir),Iy(Ic));
    end

    % Fusion
    Mh = hmxFusion(Mh);
end
end

% function Mh = hmxHorzcat(Ml,Mr)
% %+========================================================================+
% %|                                                                        |
% %|         OPENHMX - LIBRARY FOR H-MATRIX COMPRESSION AND ALGEBRA         |
% %|           openHmx is part of the GYPSILAB toolbox for Matlab           |
% %|                                                                        |
% %| COPYRIGHT : Matthieu Aussal (c) 2017-2018.                             |
% %| PROPERTY  : Centre de Mathematiques Appliquees, Ecole polytechnique,   |
% %| route de Saclay, 91128 Palaiseau, France. All rights reserved.         |
% %| LICENCE   : This program is free software, distributed in the hope that|
% %| it will be useful, but WITHOUT ANY WARRANTY. Natively, you can use,    |
% %| redistribute and/or modify it under the terms of the GNU General Public|
% %| License, as published by the Free Software Foundation (version 3 or    |
% %| later,  http://www.gnu.org/licenses). For private use, dual licencing  |
% %| is available, please contact us to activate a "pay for remove" option. |
% %| CONTACT   : matthieu.aussal@polytechnique.edu                          |
% %| WEBSITE   : www.cmap.polytechnique.fr/~aussal/gypsilab                 |
% %|                                                                        |
% %| Please acknowledge the gypsilab toolbox in programs or publications in |
% %| which you use it.                                                      |
% %|________________________________________________________________________|
% %|   '&`   |                                                              |
% %|    #    |   FILE       : hmxHorzcat.m                                  |
% %|    #    |   VERSION    : 0.40                                          |
% %|   _#_   |   AUTHOR(S)  : Matthieu Aussal                               |
% %|  ( # )  |   CREATION   : 14.03.2017                                    |
% %|  / 0 \  |   LAST MODIF : 14.03.2018                                    |
% %| ( === ) |   SYNOPSIS   : Concatenate H-Matrix along 2nd dimension      |
% %|  `---'  |                with the rule                                 |
% %|         |                H-Matrix > Full > Compr                       |
% %+========================================================================+
%
% % Check dimensions
% if ~(isa(Ml,'hmx') && isa(Mr,'hmx'))
%     error('hmxHorzcat.m : unavailable case')
% elseif (size(Ml,1)~=size(Mr,1))
%     error('hmxHorzcat.m : Dimensions of matrices being concatenated are not consistent')
% elseif norm(Ml.pos{1}-Mr.pos{1},'inf') > 1e-12
%     error('hmxHorzcat.m : unavailable case')
% end
%
% % Initialization
% Mh = hmx(Ml.pos{1},[Ml.pos{2};Mr.pos{2}],Ml.tol);
%
%
% %%% [H-Matrix & H-Matrix] -> H-Matrix
% if (Ml.typ == 0) && (Mr.typ == 0)
%     % Particles distance average
%     Yl1 = mean(Ml.chd{1}.pos{2},1);
%     Yl2 = mean(Ml.chd{2}.pos{2},1);
%     Yr1 = mean(Mr.chd{1}.pos{2},1);
%     Yr2 = mean(Mr.chd{2}.pos{2},1);
%
%     % Concatenation pairs
%     D     = [norm(Yl1-Yr1) norm(Yl1-Yr2) norm(Yl2-Yr1) norm(Yl2-Yr2)];
%     [r,d] = min(D);
%     if (d == 1) || (d == 4)
%         I = [1 1 ; 2 2 ; 3 3 ; 4 4];
%     else
%         I = [1 2 ; 2 1 ; 3 4 ; 4 3];
%     end
%
%     % Construction (recursion)
%     Mh.typ = 0;
%     for i = 1:4
%         Mh.chd{i} = hmxHorzcat(Ml.chd{I(i,1)},Mr.chd{I(i,2)});
%         Mh.row{i} = Ml.row{i};
%         Mh.col{i} = [Ml.col{I(i,1)};size(Ml,2)+Mr.col{I(i,2)}];
%     end
%
% %     % Construction (recursion)
% %     Mh.typ = 0;
% %     for i = 1:4
% %         Mh.chd{i} = hmxHorzcat(Ml.chd{i},Mr.chd{i});
% %         Mh.row{i} = Ml.row{i};
% %         Mh.col{i} = [Ml.col{i};size(Ml,2)+Mr.col{i}];
% %     end
%
%
% %     % Split for too big full leaf
% %     for i = 1:4
% %         if (Mh.chd{i}.typ == 2) && (min(size(Mh.chd{i})) > 100)
% %             tmp       = Mh.chd{i}
% %             Mh.chd{i} = hmxBuilder(tmp.pos{1},tmp.pos{2},tmp.dat,tmp.tol);
% %         end
% %     end
% %
% %     % Fusion
% %     Mh = hmxFusion(Mh);
%
% % [H-Matrix & Compr] -> H-Matrix
% elseif (Ml.typ == 0) && (Mr.typ == 1)
%     Mr = hmxSplit(Ml.row,Mr);
%     Mh = hmxHorzcat(Ml,Mr);
%
% % [H-Matrix & Full] -> H-Matrix
% elseif (Ml.typ == 0) && (Mr.typ == 2)
%     Mr = hmxSplit(Ml.row,Mr);
%     Mh = hmxHorzcat(Ml,Mr);
%
%
% %%% [Compr & H-matrix] -> H-Matrix
% elseif (Ml.typ == 1) && (Mr.typ == 0)
%     Ml = hmxSplit(Mr.row,Ml);
%     Mh = hmxHorzcat(Ml,Mr);
%
% % [Compr & Compr] -> Compr
% elseif (Ml.typ == 1) && (Mr.typ == 1)
%     % Concatenation
%     nl = size(Ml.dat{1},2);
%     nr = size(Mr.dat{1},2);
%     A  = [Ml.dat{1} , Mr.dat{1}];
%     B  = [Ml.dat{2} , zeros(nl,size(Mr,2)) ;
%         zeros(nr,size(Ml,2)) , Mr.dat{2}] ;
%
%     % Recompression
%     [A,B] = hmxQRSVD(A,B,Ml.tol);
%
%     % Update
%     Mh.typ = 1;
%     Mh.dat = {A,B};
%
%
% % [Compr & Full] -> Full
% elseif (Ml.typ == 1) && (Mr.typ == 2)
%     Mh.typ = 2;
%     Mh.dat = [Ml.dat{1}*Ml.dat{2} , Mr.dat];
%
%
% %%% [Full & H-matrix] -> H-Matrix
% elseif (Ml.typ == 2) && (Mr.typ == 0)
%     Ml = hmxSplit(Mr.row,Ml);
%     Mh = hmxHorzcat(Ml,Mr);
%
% % [Full & Compr] -> Full
% elseif (Ml.typ == 2) && (Mr.typ == 1)
%     Mh.typ = 2;
%     Mh.dat = [Ml.dat , Mr.dat{1}*Mr.dat{2}];
%
% % [Full & Full] -> Full
% elseif (Ml.typ == 2) && (Mr.typ == 2)
%     Mh.typ = 2;
%     Mh.dat = [Ml.dat,Mr.dat];
%
% else
%     error('hmxHorzcat.m : unavailable case')
% end
% end
%
%
% function Mh = hmxSplit(row,Mh)
% % Subdivision for X
% Mh.row = row;
%
% % Subdivision for Y
% [~,~,Ydim] = hmxFar(Mh);
% [I1,I2]    = hmxSubdivide(Mh.pos{2},Ydim);
% Mh.col     = {I1 , I2 , I1 , I2};
%
% % H-Matrix (recursion)
% for i = 1:4
%     % Initialization
%     Mh.chd{i} = hmx(Mh.pos{1}(Mh.row{i},:),Mh.pos{2}(Mh.col{i},:),Mh.tol);
%
%     % H-Matrix
%     if (Mh.typ == 0)
%         error('hmxSplit.m : unavailable case')
%
%     % Compressed leaf
%     elseif (Mh.typ == 1)
%         % Subdivision
%         A = Mh.dat{1}(Mh.row{i},:);
%         B = Mh.dat{2}(:,Mh.col{i});
%
%         % Recompression
%         [A,B] = hmxQRSVD(A,B,Mh.tol);
%
%         % Update
%         Mh.chd{i}.typ = 1;
%         Mh.chd{i}.dat = {A,B};
%         Mh.chd{i}.tol = Mh.tol;
%
%     % Full leaf
%     elseif (Mh.typ == 2)
%         Mh.chd{i}.typ = 2;
%         Mh.chd{i}.dat = Mh.dat(Mh.row{i},Mh.col{i});
%         Mh.chd{i}.tol = Mh.tol;
%     end
% end
%
% % Erase data
% Mh.typ = 0;
% Mh.dat = [];
% end


function [CTTraceS, RTOGPBVectorsM, RTOGPBVectorsM_MC, PBMaskM, rowPBV, colPBV, xPBPosV, yPBPosV, beamlet_delta_x, beamlet_delta_y] = ...
   getPBRayData_cerr3(sourceS, numSamplePts, xySampleRate,planC, xPosV, yPosV, beamlet_delta_x, beamlet_delta_y, gA);
% function [CTTraceS, RTOGPBVectorsM, RTOGPBVectorsM_MC, PBMaskM, rowPBV, colPBV, xPBPosV, yPBPosV, beamlet_delta_x, beamlet_delta_y] = ...
%    getPBRayData_cerr3(edgeS, sourceS, numSamplePts, xySampleRate,planC, xPosV, yPosV, beamlet_delta_x, beamlet_delta_y, gA);
%JOD
% global planC
%
% indexS=planC{end};

water = 1000; %Assumes water equals 1000.

%-----------Get CT scan---------------------%

% [CTUniform3D, CTUniformInfoS] = getUniformizedCTScan(1,planC);
scanNum = 1;
[CTUniform3D, CTUniformInfoS] = getUniformizedCTScan(1,scanNum);

xOffset = CTUniformInfoS.xOffset;
yOffset = CTUniformInfoS.yOffset;

%-----------Fix source characteristics---------------------%

orgV   = [sourceS.isocenter.x, sourceS.isocenter.y, sourceS.isocenter.z];

%-----------Get ray parameters---------------------%

% [RTOGPBVectorsM, RTOGPBVectorsM_MC, PBMaskM, rowPBV, colPBV, xPBPosV, yPBPosV] = ...
%       getPBRays(edgeS, sourceS, xySampleRate);
% [RTOGPBVectorsM, RTOGPBVectorsM_MC, PBMaskM, rowPBV, colPBV, xPBPosV, yPBPosV, beamlet_delta_x, beamlet_delta_y] = ...
%       getPBRays(xPosV, yPosV, beamlet_delta_x, beamlet_delta_y, gA);
[RTOGPBVectorsM, RTOGPBVectorsM_MC, PBMaskM, rowPBV, colPBV, xPBPosV, yPBPosV, beamlet_delta_x, beamlet_delta_y] = ...
      getPBRays_planChk(xPosV, yPosV, beamlet_delta_x, beamlet_delta_y, sourceS);

rayLength = 500;  %in cm.  This is the length of the rays which are passed
                  %through the CT matrix to determine cumulative CT densities.

numSlices = size(CTUniform3D,3);

zFirst = CTUniformInfoS.firstZValue;

sliceThickness = CTUniformInfoS.sliceThickness;

%minBox an maxBox containt coordinates of the corners of the CT box as used
%by the ray intersection routine.

maxBoxS.z = (numSlices - 1) * sliceThickness + zFirst + sliceThickness/2;

minBoxS.z = zFirst - sliceThickness/2;

delta_xy = CTUniformInfoS.grid1Units;

imageWidth = CTUniformInfoS.size(1);

minBoxS.x = - imageWidth/2 * delta_xy + xOffset;
maxBoxS.x =   imageWidth/2 * delta_xy + xOffset;

minBoxS.y = - imageWidth/2 * delta_xy + yOffset;
maxBoxS.y =   imageWidth/2 * delta_xy + yOffset;

imageWidth = CTUniformInfoS.sizeOfDimension1;

CTTraceS = struct('CTNumsRay',[],'CTCumNumsRay',[],'distSamplePts',[]);

for i = 1 : size(RTOGPBVectorsM,1)

  rayDeltaS.x = RTOGPBVectorsM(i,1) * rayLength;
  rayDeltaS.y = RTOGPBVectorsM(i,2) * rayLength;
  rayDeltaS.z = RTOGPBVectorsM(i,3) * rayLength;

  %are the components of the ray's direction and maximum length.

  deltaV = [rayDeltaS.x, rayDeltaS.y, rayDeltaS.z];

  t_entrance = rayBoxIntersection(sourceS,rayDeltaS,minBoxS,maxBoxS);

  %The entrance point is therefore
  entranceV = orgV + t_entrance * deltaV;

  %Reflect to find exit point (assume length of ray is long enough that ray does exit):

  if t_entrance ~= -1

    %find exit point
    %get end of ray
    rayOrgS2.xRel = sourceS.xRel + rayDeltaS.x;  %reflected source positions
    rayOrgS2.yRel = sourceS.yRel + rayDeltaS.y;
    rayOrgS2.zRel = sourceS.zRel + rayDeltaS.z;

    rayDeltaS2.x = - rayDeltaS.x;
    rayDeltaS2.y = - rayDeltaS.y;
    rayDeltaS2.z = - rayDeltaS.z;

    rayOrgS2.isocenter = sourceS.isocenter;

    t = rayBoxIntersection(rayOrgS2,rayDeltaS2,minBoxS,maxBoxS);

    t_exit = 1 - t;

    exitV = orgV + t_exit * deltaV;

else
    error('PB Ray does not intersect CT scan.');
end

  %Now produce a set of sampling points between the entrance and the exit

  nV = 1 : numSamplePts;

  delta_t  = (t_exit - t_entrance)/(numSamplePts - 1);

  tV = t_entrance + (nV - 1) * delta_t;

  CTTraceS(i).distSamplePts = tV * sum(deltaV.^2).^0.5;

  sampleV.x =  sourceS.xRel + tV * deltaV(1);

  sampleV.y =  sourceS.yRel + tV * deltaV(2);

  sampleV.z =  sourceS.zRel + tV * deltaV(3);

  sampleRTOGV.x = sampleV.x + sourceS.isocenter.x;
  sampleRTOGV.y = sampleV.y + sourceS.isocenter.y;
  sampleRTOGV.z = sampleV.z + sourceS.isocenter.z;

  %---------Sample CT densities----------%

  %To go from sample points in RTOG system to CT densities, we convert as follows:

  %What is the slice number?
  sliceV = 1 + (sampleRTOGV.z - zFirst)/sliceThickness;

  %Now do 3-D interpolation:
  zFieldV = [minBoxS.z + 0.5 * sliceThickness:sliceThickness:maxBoxS.z - 0.5 * sliceThickness];
  xFieldV = [minBoxS.x + 0.5 * delta_xy, delta_xy, maxBoxS.x - 0.5 * delta_xy];
  yFieldV = [minBoxS.y + 0.5 * delta_xy, delta_xy, maxBoxS.y - 0.5 * delta_xy];

  [CTNumsV] = finterp3(sampleRTOGV.x, sampleRTOGV.y, sampleRTOGV.z, CTUniform3D, xFieldV, yFieldV, zFieldV, 0);

%   [CTNumsV] = finterp3(sampleRTOGV.x, sampleRTOGV.y, sampleRTOGV.z, CTUniform3D, xFieldV, yFieldV, zFieldV, 0);

  CTTraceS(i).densityRay = delta_t * norm(deltaV) * CTNumsV/water;
  CTTraceS(i).cumDensityRay = cumsum(CTTraceS(i).densityRay);  %Account for sampling rate to convert to g/cm^2.

  if any(CTTraceS(i).cumDensityRay > 50)
      warning('Cumulative density ray appears to exceed maximum length.');
  end

end

%-----------fini---------------------%


function z = times( x, y, oper )

%   Disciplined convex programming information for TIMES:
%      In general, disciplined convex programs must obey the "no-product
%      rule" which states that two non-constant expressions cannot be
%      multiplied together. Under this rule, the following combinations
%      are allowed:
%         {constant} .* {non-constant}  {non-constant} .* {constant}
%      Furthermore, if the non-constant expression is nonlinear, then
%      the constant term must be real.
%
%      A lone exception to the no-product rule is made for quadratic
%      forms: two affine expressions may be multiplied together if the
%      result is convex or concave. For example, the construction
%         variable x(n)
%         x.*x  <= 1;
%      would be permitted because each element of x.*x is convex.
%
%   Disciplined geometric programming information for TIMES:
%      Both terms in a multiplication must have the same log-curvature,
%      so the following products are permitted:
%         {log-convex} .* {log-convex}  {log-concave} .* {log-concave}
%         {log-affine} .* {log-affine}
%      Note that log-affine expressions are both log-convex and
%      log-concave.
%
%   For vectors, matrices, and arrays, these rules are verified
%   indepdently for each element.

persistent P
if isempty( P ),
    P.map = cvx_remap( ...
        ... % Invalid combinations: negative of log-concave, gp term
        { { 'negative' },   { 'l_concave_', 'g_valid' } }, ...
        { { 'l_concave_', 'g_valid' }, { 'negative' }   }, ...
        ... % Left-multiply by constant
        { { 'constant' },   { 'affine' }     }, ...
        { { 'real' },       { 'valid'  }     }, ...
        ... % Right-multiply by constant
        { { 'affine' },     { 'constant' }   }, ...
        { { 'valid' },      { 'real'  }      }, ...
        ... % Geometric
        { { 'l_convex' } }, ...
        { { 'l_concave' } }, ...
        ... % Potential quadratic form
        { { 'affine', 'p_convex', 'n_concave' } }, ...
        [ 0, 0, 1, 1, 2, 2, 3, 3, 4 ] );
    P.funcs = { @times_l, @times_r, @times_g, @times_q };
    P.constant = [];
end
if nargin < 3, oper = '.*'; end
P.name = oper;
z = cvx_binary_op( P, x, y );

function z = times_l( x, y )
% constant .* something
z = bsxfun( @times, cvx_constant( x ).', cvx_basis( y ) );
z = cvx( size(z,2), z );

function z = times_r( x, y )
% constant .* something
z = bsxfun( @times, cvx_basis( x ), cvx_constant( y ).' );
z = cvx( size(z,2), z );

function z = times_g( x, y )
% geom .* geom
z = exp( plus_nc( log( x ), log( y ) ) );

function z = times_q( x, y )
% affine .* affine
% p_convex .* p_convex
% n_concave .* n_concave
xA = x.basis_;
yA = y.basis_;
mm = max( size( xA, 1 ), size( yA, 1 ) );
xA( end + 1 : mm, : ) = 0;
yA( end + 1 : mm, : ) = 0;
xB  = xA( 2 : end, : );
yB  = yA( 2 : end, : );
cyB = conj( yB );
alpha = sum( yB .* cyB, 1 );
alpha = sum( bsxfun( @times, xB, yB ), 1 ) ./ max( alpha, realmin );
if nnz( xB - bsxfun( @times, alpha, cyB ) > 2 * eps * sqrt( conj( xB ) .* xB ) ),
    cvx_throw( 'Disciplined convex programming error:\n    Invalid quadratic form(s): not a square.' );
elseif any( abs(imag(alpha)) > 2 * eps * abs(real(alpha)) ),
    cvx_throw( 'Disciplined convex programming error:\n    Invalid quadratic form(s): not real.' );
else
    if isreal( xB ),
        z = square( y );
    else
        z = square( abs( y ) );
    end
    z.basis_ = bsxfun( @times, alpha, z.basis_ );
    offset = xA(1,:) - alpha .* conj( yA(1,:) );
    if any( offset ),
        z = z + offset.' .* y;
    end
end

% Copyright 2005-2014 CVX Research, Inc.
% See the file LICENSE.tx for full copyright information.
% The command 'cvx_where' will show where this file is located.


function Lxx = nlp_hessfcn(om, x, lambda, cost_mult, Hs)
%NLP_HESSFCN  Evaluates Hessian of Lagrangian.
%   LXX = NLP_HESSFCN(OM, X, LAMBDA, COST_MULT)
%   LXX = NLP_HESSFCN(OM, X, LAMBDA, COST_MULT, HS)
%
%   Hessian evaluation function, suitable for use with MIPS, FMINCON, etc.
%
%   Inputs:
%     OM : Opt-Model object
%     X : optimization vector
%     LAMBDA (struct)
%       .eqnonlin : Lagrange multipliers on equality constraints
%       .ineqnonlin : Kuhn-Tucker multipliers on inequality constraints
%     COST_MULT : (optional) Scale factor to be applied to the cost
%          (default = 1).
%     HS : (optional) sparse matrix with tiny non-zero values specifying
%          the fixed sparsity structure that the resulting LXX should match
%
%   Outputs:
%     LXX : Hessian of the Lagrangian.
%
%   Examples:
%       Lxx = nlp_hessfcn(om, x, lambda, cost_mult);
%       Lxx = nlp_hessfcn(om, x, lambda, cost_mult, Hs);
%
%   See also NLP_COSTFCN, NLP_CONSFCN.

%   MP-Opt-Model
%   Copyright (c) 1996-2020, Power Systems Engineering Research Center (PSERC)
%   by Ray Zimmerman, PSERC Cornell
%
%   This file is part of MP-Opt-Model.
%   Covered by the 3-clause BSD License (see LICENSE file for details).
%   See https://github.com/MATPOWER/mp-opt-model for more info.

%% ----- evaluate d2f -----
[f, df, d2f] = nlp_costfcn(om, x);
d2f = d2f * cost_mult;

%%----- evaluate Hessian of equality constraints -----
d2G = om.eval_nln_constraint_hess(x, lambda.eqnonlin, 1);

%%----- evaluate Hessian of inequality constraints -----
d2H = om.eval_nln_constraint_hess(x, lambda.ineqnonlin, 0);

%%-----  do numerical check using (central) finite differences  -----
if 0
    nx = length(x);
    step = 1e-5;
    num_d2f = sparse(nx, nx);
    num_d2G = sparse(nx, nx);
    num_d2H = sparse(nx, nx);
    for i = 1:nx
        xp = x;
        xm = x;
        xp(i) = x(i) + step/2;
        xm(i) = x(i) - step/2;
        % evaluate cost & gradients
        [fp, dfp] = nlp_costfcn(om, xp);
        [fm, dfm] = nlp_costfcn(om, xm);
        % evaluate constraints & gradients
        [Hp, Gp, dHp, dGp] = nlp_consfcn(om, xp);
        [Hm, Gm, dHm, dGm] = nlp_consfcn(om, xm);
        num_d2f(:, i) = cost_mult * (dfp - dfm) / step;
        num_d2G(:, i) = (dGp - dGm) * lambda.eqnonlin   / step;
        num_d2H(:, i) = (dHp - dHm) * lambda.ineqnonlin / step;
    end
    d2f_err = full(max(max(abs(d2f - num_d2f))));
    d2G_err = full(max(max(abs(d2G - num_d2G))));
    d2H_err = full(max(max(abs(d2H - num_d2H))));
    if d2f_err > 1e-6
        fprintf('Max difference in d2f: %g\n', d2f_err);
    end
    if d2G_err > 1e-5
        fprintf('Max difference in d2G: %g\n', d2G_err);
    end
    if d2H_err > 1e-6
        fprintf('Max difference in d2H: %g\n', d2H_err);
    end
end

Lxx = d2f + d2G + d2H;

%% force specified sparsity structure
if nargin > 4
    %% add sparse structure (with tiny values) to current matrices to
    %% ensure that sparsity structure matches that supplied
    Lxx = Lxx + Hs;

%     %% check sparsity structure against that supplied
%     if nnz(Lxx) ~= nnz(Hs)
%         fprintf('=====> nnz(Lxx) is %d, expected %d <=====\n', nnz(Lxx), nnz(Hs));
%     else
%         [iHs, jHs] = find(Hs);
%         [iH, jH] = find(Lxx);
%         if any(iH ~= iHs) || any(jH ~= jHs)
%             fprintf('=====> structure of Lxx is not as expected <=====\n');
%         end
%     end
end


function x = vec(X)
% x = vec(X)
%
% Y = VEC(x)  Given an m x n matrix x, this produces the vector Y of length
%   m*n that contains the columns of the matrix x, stacked below each other.
%
% See also mat.

% This file is part of SeDuMi 1.1 by Imre Polik and Oleksandr Romanko
% Copyright (C) 2005 McMaster University, Hamilton, CANADA  (since 1.1)
%
% Copyright (C) 2001 Jos F. Sturm (up to 1.05R5)
%   Dept. Econometrics & O.R., Tilburg University, the Netherlands.
%   Supported by the Netherlands Organization for Scientific Research (NWO).
%
% Affiliation SeDuMi 1.03 and 1.04Beta (2000):
%   Dept. Quantitative Economics, Maastricht University, the Netherlands.
%
% Affiliations up to SeDuMi 1.02 (AUG1998):
%   CRL, McMaster University, Canada.
%   Supported by the Netherlands Organization for Scientific Research (NWO).
%
% This program is free software; you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation; either version 2 of the License, or
% (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program; if not, write to the Free Software
% Foundation, Inc.,  51 Franklin Street, Fifth Floor, Boston, MA
% 02110-1301, USA

x = reshape(X,numel(X),1);


function value = r8vec_in_ab ( n, x, a, b )

%*****************************************************************************80
%
%% R8VEC_IN_AB is TRUE if the entries of an R8VEC are in the range [A,B].
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    26 March 2012
%
%  Author:
%
%    John Burkardt
%
%  Parameters:
%
%    Input, integer N, the number of entries in the vector.
%
%    Input, real X(N), the vector.
%
%    Input, real A, B, the limits, with A <= B.
%
%    Output, logical VALUE, is TRUE if every entry is
%    between A and B.
%
  if ( any ( x(1:n) < a | b < x(1:n) ) )
    value = 0;
  else
    value = 1;
  end

  return
end


function c = diag(a,k)
%DIAG         Implements  diag(a,k)  for hessians
%
%   c = diag(a,k)
%
% functionality as Matlab function diag for matrices
%

% written  04/04/04     S.M. Rump
% modified 04/06/05     S.M. Rump  rounding unchanged
% modified 08/26/12     S.M. Rump  global variables removed
%

  INTLAB_HESSIAN_NUMVAR = getappdata(0,'INTLAB_HESSIAN_NUMVAR');

  if nargin==1
    k = 0;
  end

  [m n] = size(a.x);
  c.x = diag(a.x,k);
  if m==1 | n==1
    index = diag(ones(1,m*n),k);
    if issparse(a.hx)
      c.dx = sparse([],[],[],INTLAB_HESSIAN_NUMVAR,prod(size(index)),0);
      c.hx = sparse([],[],[],INTLAB_HESSIAN_NUMVAR^2,prod(size(index)),0);
    else
      c.dx = zeros(INTLAB_HESSIAN_NUMVAR,prod(size(index)));
      c.hx = zeros(INTLAB_HESSIAN_NUMVAR^2,prod(size(index)));
    end
    if isa(a.x,'intval')
      c.dx = intval(c.dx);
      c.hx = intval(c.hx);
    end
    c.dx(:,index~=0) = a.dx;
    c.hx(:,index~=0) = a.hx;
  else
    index = diag( reshape( 1:(m*n) , m , n ) , k );
    c.dx = a.dx(:,index(:));
    c.hx = a.hx(:,index(:));
  end

  c = class(c,'hessian');


%% FUNCTION Least_Dirty
%   Dirty Multi-Task Learning with Least Squares Loss.
%
%% OBJECTIVE
%   argmin_W ||X(P+Q) - Y||_F^2 + lambda1*||P||_{1,inf} + lambda2*||Q||_{1,1}
%    s.t. W = P + Q
%
%% INPUT
%   X: {n * d} * t - input matrix
%   Y: {n * 1} * t - output matrix
%   rho1: group sparsity regularization parameter
%   rho2: elementwise sparsity regularization parameter
%
%% OUTPUT
%   W: model: d * t
%   funcVal: function value vector.
%   P: group sparsity structure (joint feature slection)
%   Q: elementwise sparsity component
%
%% LICENSE
%   This program is free software: you can redistribute it and/or modify
%   it under the terms of the GNU General Public License as published by
%   the Free Software Foundation, either version 3 of the License, or
%   (at your option) any later version.
%
%   This program is distributed in the hope that it will be useful,
%   but WITHOUT ANY WARRANTY; without even the implied warranty of
%   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%   GNU General Public License for more details.
%
%   You should have received a copy of the GNU General Public License
%   along with this program.  If not, see <http://www.gnu.org/licenses/>.
%
%   Copyright (C) 2011 - 2012 Jiayu Zhou, Pinghua Gong and Jieping Ye
%
%   You are suggested to first read the Manual.
%   For any problem, please contact with Jiayu Zhou via jiayu.zhou@asu.edu
%
%   Last modified on June 3, 2012.
%
%% RELATED PAPERS
%
%   [1] Jalali, A. and Ravikumar, P. and Sanghavi, S. and Ruan, C. A dirty
%       model for multi-task learning, NIPS 2010.
%
%% RELATED FUNCTIONS
%   init_opts, combine_input (utils), prf_lbm (c_files)

function [W, funcVal, P, Q] = Least_Dirty(X, Y, lambda1, lambda2, opts)

if nargin <4
    error('\n Inputs: X, Y, and lambda1, and lambda2 should be specified!\n');
end
if nargin <5
    opts = [];
end

% initialize options.
opts=init_opts(opts);

% initial Lipschiz constant.
if isfield(opts, 'lFlag')
    lFlag = opts.lFlag;
else
    lFlag = false;
end

task_num = length(X);
[X, y, ~, samplesize] = combine_input(X, Y);
dimension = size(X, 2);

% initialize a starting point
if opts.init==2
    P0 = zeros(dimension, task_num);
    Q0 = zeros(dimension, task_num);
elseif opts.init == 0
    P0 = randn(dimension, task_num);
    Q0 = randn(dimension, task_num);
else
    if isfield(opts,'P0')
        P0=opts.P0;
        if (nnz(size(P0)-[dimension, task_num]))
            error('\n Check the input .P0');
        end
    else
        P0=zeros(dimension, task_num);
    end

    if isfield(opts,'Q0')
        Q0=opts.Q0;
        if (nnz(size(Q0)-[dimension, task_num]))
            error('\n Check the input .Q0');
        end
    else
        Q0=zeros(dimension, task_num);
    end
end

% Set an array to save the objective value
funcVal = [];

P = P0; Q = Q0;

%[d,m] = size(P); % d: dimension, m: the number of tasks
X = diagonalize(X,samplesize);
XtX = X'*X; Xty = X'*y;

Pn = P;  Qn = Q;
t_new = 1;
L1norm = max(sum(abs(X),1)); Linfnorm = max(sum(abs(X),2));

if lFlag
    % Upper bound for largest eigenvalue of Hessian matrix
    L = 2*min([L1norm*Linfnorm; size(X,1)*Linfnorm*Linfnorm; size(X,2)*L1norm*L1norm; size(X,1)*size(X,2)*max(abs(X(:)))]);
else
    % Lower bound for largest eigenvalue of Hessian matrix
    L = 2*max(L1norm*L1norm/size(X,1),Linfnorm*Linfnorm/size(X,2));
end
% Initial function value
funcVal = cat(1, funcVal, norm(X*(P(:) + Q(:)) - y)^2 + lambda1*L1infnorm(P) + lambda2*L11norm(Q));
%fun(1) = norm(X*(P(:) + Q(:)) - y)^2 + lambda1*L1infnorm(P) + lambda2*L11norm(Q);

%count = 0;
for iter = 1 : opts.maxIter
    P_old = P; Q_old = Q;
    t_old = t_new;
    gradvec = 2*(XtX*(Pn(:)+Qn(:)) - Xty);
    gradmat = reshape(gradvec,dimension,task_num);
    % If we estimate the upper bound of Lipschitz constant, no line search
    % is needed.
    if lFlag
        P = proximalL1infnorm(Pn - gradmat/L, lambda1/L);
        Q = proximalL11norm(Qn - gradmat/L, lambda2/L);
    else
        % line search
        for inneriter = 1:20
            P = proximalL1infnorm(Pn - gradmat/L, lambda1/L);
            Q = proximalL11norm(Qn - gradmat/L, lambda2/L);
            dP = P - Pn;  dQ = Q - Qn;
            if 2*((dP(:) + dQ(:))'*XtX*(dP(:) + dQ(:))) <= L*sum(sum((dP.*dP + dQ.*dQ)))
                break;
            else
                L = L*2;
            end
        end
    end
    funcVal = cat(1, funcVal, norm(X*(P(:) + Q(:)) - y)^2 + lambda1*L1infnorm(P) + lambda2*L11norm(Q));
    %fun(iter+1) = norm(X*(P(:) + Q(:)) - y)^2 + lambda1*L1infnorm(P) + lambda2*L11norm(Q);

    % test stop condition.
    switch(opts.tFlag)
        case 0
            if iter>=2
                if (abs( funcVal(end) - funcVal(end-1) ) <= opts.tol)
                    break;
                end
            end
        case 1
            if iter>=2
                if (abs( funcVal(end) - funcVal(end-1) ) <=...
                        opts.tol* funcVal(end-1))
                    break;
                end
            end
        case 2
            if ( funcVal(end)<= opts.tol)
                break;
            end
        case 3
            if iter>=opts.maxIter
                break;
            end
    end

    %     % stopping condition
    %     if abs(fun(iter) - fun(iter+1))/fun(iter+1) < tol
    %         count = count + 1;
    %     else
    %         count = 0;
    %     end
    %     if count >= 3
    %         break;
    %     end

    % Update the coefficient
    t_new = (1+sqrt(1+4*t_old^2))/2;
    Pn = P + (t_old-1)/t_new*(P - P_old);
    Qn = Q + (t_old-1)/t_new*(Q - Q_old);

end

W = P + Q;

end


function ADEM_cost_SHC
% This demo illustrates the use of priors on the motion of hidden states as
% polices. It simulates exploration and exploitation using radial basis
% function attractors and auto-vitiative (self-destroying) attractors
% as the basis of the prior. These dynamics enforce exploration, under
% active inference. In turn, this exploration enables perceptual learning
% to associate attractors with changes in physiological states (cf,
% rewards). This can be exploited to by formal priors to ensure regions of
% physiological state-space are avoided.
% We look at this scheme using simulated pathology; first, we simulate a
% (neurodegenerative) reduction in log-precision (cf Parkinson's disease) on
% the motion of physical states.  This results in active inference with
% a loss of precise volitional movement and subsequent failure to optimise
% physiological states. Second, we look at the effects of precision on
% learning by increasing log-precision (cf, Addition) on the motion of
% physiological states. This results in a failure to learn and, again,
% subsequent failure to optimise physiological states.
%__________________________________________________________________________
% Copyright (C) 2010 Wellcome Trust Centre for Neuroimaging

% Karl Friston
% $Id: ADEM_cost_SHC.m 7679 2019-10-24 15:54:07Z spm $

% switch for demo
%--------------------------------------------------------------------------
DEMO = 1;

% location and radius of attractors (A)
%--------------------------------------------------------------------------
global A
A.x  = [1 -1 -1  1;
        2 -2  2 -2];
A.d  = 1/4;              % s.d. of Gaussian radial basis function
A.u  = 1/8;              % threshold to induce rapid decay of x.a
A.q  = [1 2];            % indices of locations that increase x.q

% generative process
%==========================================================================

% level 1
%--------------------------------------------------------------------------
G(1).x.x = [0;0];
G(1).x.v = [1;0];
G(1).x.q = [1;1/2];

G(1).f   = inline('spm_cost_SHC_fxa(x,v,a,P)','x','v','a','P');
G(1).g   = inline('[x.x; x.v; x.q]','x','v','a','P');
G(1).V   = exp(16);                          % error precision
G(1).W   = exp(16);                          % error precision
G(1).U   = exp(4);                           % action precision

% level 2
%--------------------------------------------------------------------------
G(2).a   = sparse(2,1);                      % action
G(2).v   = 0;                                % inputs
G(2).V   = exp(16);
G        = spm_ADEM_M_set(G);

% generative model
%==========================================================================

% level 1
%--------------------------------------------------------------------------
M(1).x.x = G(1).x.x;
M(1).x.v = G(1).x.v;
M(1).x.q = G(1).x.q;
M(1).x.a = randn(4,1)/8;

M(1).f   = inline('spm_cost_SHC_fx(x,v,P)','x','v','P');
M(1).g   = inline('[x.x; x.v; x.q]','x','v','P');
M(1).pE  = speye(4,2);

M(1).V   = exp(8);
M(1).W   = diag(exp([[1 1 1 1]*4 [1 1]*6 [1 1 1 1]*4]));

% level 2 (no exogenous inputs in this simulation)
%--------------------------------------------------------------------------
M(2).v   = 0;
M(2).V   = exp(16);

% Integrate: active inference
%==========================================================================
M(1).E.nE = 1;
M(1).E.n  = 4;

N      = 128;
DEM.U  = sparse(N,1);
DEM.C  = sparse(N,1);
DEM.G  = G;
DEM.M  = M;

if DEMO
    load DEM_addiction
else
    DEM = spm_ADEM(DEM);
end

% show behavior
%==========================================================================

% overview
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 1');
spm_DEM_qU(DEM.qU)

subplot(2,2,3)
spm_cost_SHC_path(DEM.pU,A)

% a closer look at physiology
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 2'); clf

subplot(2,1,1)
t    = 1:N;
plot(t,DEM.qU.x{1}(5:6,:)',t,t*0 + A.u,':'), hold on
plot(t,DEM.qU.x{1}(7:10,:)','-.'),           hold off
xlabel('time','Fontsize',14)
title('internal states','FontSize',16)
axis square, box off, set(gca,'XLim',[1 N])

subplot(2,2,3)
plot(t,DEM.qU.x{1}(5:6,:)',t,t*0 + A.u,':')
xlabel('time','Fontsize',14)
title('physiological states','FontSize',16)
axis square, box off, set(gca,'XLim',[1 N])

subplot(2,2,4)
plot(DEM.pU.x{1}(5,:),DEM.pU.x{1}(6,:)),               hold on
plot([A.u 1],[A.u A.u],'-.r',[A.u A.u],[A.u 1],'-.r'), hold off
xlabel('time','Fontsize',14)
title('physiological states','FontSize',16)
axis square, box off, axis([-.1 1.2 -.1 1.2])

% look at the effect of precision on inference (cf Parkinson's disease)
%==========================================================================

% initialize states
%--------------------------------------------------------------------------
DEM = spm_ADEM_update(DEM);

% simulate exposure with different log-precisions on physical motion
%--------------------------------------------------------------------------
WP    = [4 2 0];
if ~DEMO
    for i = 1:length(WP)
        DEM_P{i}        = DEM;
        W               = [[1 1 1 1]*WP(i) [1 1]*6 [1 1 1 1]*4];
        DEM_P{i}.M(1).W = diag(exp(W));
        DEM_P{i}        = spm_ADEM(DEM_P{i});
    end
end

% Graphics - path and physiology
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 3'); clf
for i = 1:3

    % path
    %----------------------------------------------------------------------
    spm_figure('GetWin','Figure 3');
    subplot(3,2,i*2 - 1)
    spm_cost_SHC_path(DEM_P{i}.pU,A)

    % and physiology
    %----------------------------------------------------------------------
    subplot(3,2,i*2)
    plot(t,DEM_P{i}.pU.x{1}(5:6,:)',t,t*0 + A.u,':'), hold on
    plot(t,sum(DEM_P{i}.pU.x{1}(1:2,:).^2)/8,'m:'),    hold off
    xlabel('time','Fontsize',12)
    title(sprintf('%s (%1.0f)','physiological states',WP(i)),'FontSize',16)
    axis square, box off, set(gca,'XLim',[1 N])
    drawnow

end

% Graphics - action and prediction error
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 4'); clf
for i = 1:3

    % path
    %----------------------------------------------------------------------
    subplot(3,2,2*i - 1)
    plot(t,DEM_P{i}.qU.a{2})
    xlabel('time','FontSize',12)
    title('action','FontSize',16)
    axis square, box off, set(gca,'XLim',[1 N])
    if i == 1; a = axis; end, axis(a)

    % and prediction error
    %----------------------------------------------------------------------
    subplot(3,2,2*i)
    plot(t,DEM_P{i}.qU.z{1}(1:4,:))
    xlabel('time','FontSize',12)
    title('sensory error','FontSize',16)
    axis square, box off, set(gca,'XLim',[1 N])
    if i == 1; aa = axis; end, axis(aa)
    drawnow
end

% Learning (at different levels of precision on physiological motion)
%==========================================================================

% Switch attractor for second physiological state and log-precision levels
%--------------------------------------------------------------------------
spm_figure('GetWin','DEM'); clf
A.q  = [1 3];
WA   = [4 8 12];

if ~DEMO

    for i = 1:length(WA)

        % Enable learning
        %------------------------------------------------------------------
        A.u              = 0;
        DEM_L{i}         = DEM;
        DEM_L{i}.M(1).pC = exp(8);
        W                = [[1 1 1 1]*4 [1 1]*WA(i) [1 1 1 1]*4];
        DEM_L{i}.M(1).W  = diag(exp(W));
        DEM_L{i}         = spm_ADEM(DEM_L{i});

        % replace prior with posterior and re-expose
        %------------------------------------------------------------------
        A.u              = 1/8;
        DEM_D{i}         = DEM_L{i};
        DEM_D{i}.M(1).pE = DEM_D{i}.qP.P{1};
        DEM_D{i}.M(1).pC = [];
        DEM_D{i}         = spm_ADEM(DEM_D{i});

    end

    % save DEM structures for future demonstrations
    %----------------------------------------------------------------------
    save DEM_addiction DEM DEM_P DEM_L DEM_D

end

% Graphics - optimal learning
%--------------------------------------------------------------------------
spm_figure('GetWin','DEM'); clf
spm_DEM_qU(DEM.qU,DEM.pU)

spm_figure('GetWin','Figure 5'); clf
spm_DEM_qP(DEM_L{1}.qP)

subplot(2,2,3)
spm_cost_SHC_path(DEM.pU,A)
title('Before','Fontsize',16)
subplot(2,2,4)
spm_cost_SHC_path(DEM_D{1}.pU,A)
title('After','Fontsize',16)
drawnow

% Comparison of learning over levels of log-precisions (cf. Addiction)
%==========================================================================
spm_figure('GetWin','Figure 6'); clf

for i = 1:3

    % path
    %----------------------------------------------------------------------
    subplot(3,2,i*2 - 1)
    spm_cost_SHC_path(DEM_D{i}.pU,A)
    title('After','FontSize',16)

    % and physiology
    %----------------------------------------------------------------------
    subplot(3,2,i*2)
    spm_plot_ci(DEM_L{i}.qP.P{1}(:),DEM_L{i}.qP.C)
    xlabel('parameter','FontSize',12)
    title(sprintf('%s (%1.0f)','learning',WA(i)),'FontSize',16)
    axis square, box off

end
drawnow

% and physiology
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 7'); clf
for i = 1:3

    % predicted motion of second physiological state
    %----------------------------------------------------------------------
    for j = 1:N
        x    = DEM_L{i}.qU.x{1}(:,j);
        f    = spm_cost_SHC_fx(spm_unvec(x,M(1).x),M(2).v,M(1).pE);
        p(j) = DEM_L{i}.qU.w{1}(6,j) + f.q(2);
        x    = DEM_L{i}.pU.x{1}(:,j);
        f    = spm_cost_SHC_fxa(spm_unvec(x,G(1).x),M(2).v,G(2).a,G(1).pE);
        q(j) = f.q(2);
    end
    subplot(3,2,2*i - 1)
    plot(t,p,t,q,'b:')
    xlabel('time','FontSize',12)
    title(sprintf('%s (%1.0f)','predicted motion',WA(i)),'FontSize',16)
    axis square, box off, set(gca,'XLim',[1 N])
    if i == 1; a = axis; end, axis(a)

    subplot(3,2,2*i)
    plot(t,DEM_L{i}.qU.w{1}(6,:))
    xlabel('time','FontSize',12)
    title('prediction error','FontSize',16)
    axis square, box off, set(gca,'XLim',[1 N])
    if i == 1; aa = axis; end, axis(aa)

end


function [IPPEPoses,refinedPoses] = perspectiveIPPE(U,Q,hEstMethod,opts)
%perspectiveIPPE: The solution to Perspective IPPE with point correspondences computed
%between points in world coordinates on the plane z=0, and normalised points in the
%camera's image.
%
%Inputs:
%
%U: 2xN or 3xN matrix holding the model points in world coordinates. If U
%is 2xN then the points are defined in world coordinates on the plane z=0
%
%Q: 2xN matrix holding the points in the image. These are in normalised
%pixel coordinates. That is, the effects of the camera's intrinsic matrix
%and lens distortion are corrected, so that the Q projects with a perfect
%pinhole model.
%
%hEstMethod: the homography estimation method, either 'Harker' (default) or 'DLT'. If
%'Harker' then the method of Harker and O'Leary is used, from the paper
%"Computation of Homographies" in the 2005 British Computer Vision
%Conference. Otherwise the Direct Linear Transform is used, from Peter Kovesi's implementation at http://www.csse.uwa.edu.au/~pk.
%If hEstMethod=[] then the default is used.
%
%opts is an optional datastructure holding the fields:
%setup IPPE options:
%opts.withPoseRefinement = true; %If this is set, we also refine the camera pose from IPPE by gradient-based optimisation (Levenberg–Marquardt). This is designed to give the statistically optimal pose assuming zero mean I.I.D nose for the feature's 2D positions, but requires initialisation (from IPPE). Note that it about two orders of magnitude slower than IPPE.
%opts.displayTiming = true; %set this to true if you want to see how long IPPE takes in Matlab.
%
%Outputs:
%IPPEPoses: structure containing the two pose solutions from IPPE.
%   IPPEPoses.R1 is the first pose's rotation
%   IPPEPoses.t1 is the first pose's translation

%   IPPEPoses.R2 is the second pose's rotation
%   IPPEPoses.t2 is the second pose's translation

%   IPPEPoses.reprojError1 is the first pose's reprojection error
%   IPPEPoses.reprojError2 is the second pose's reprojection error
%
%refinedPoses: structure containing the two refined pose solutions using
%Levenberg–Marquardt (same format as IPPEPoses). This is only outputted if opts.withPoseRefinement = true
%
%Note that the poses are transforms from world
%coordinates to camera coordinates. If you
%want to get the camera's pose in world coordinates you would get it with
%M1 = inv([[R1,t1];[0,0,0,1]]) and M2 = inv([[R2,t2];[0,0,0,1]])
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% This file is part of the IPPE package for very fast plane-based pose
% estimation from a single perspective image, from the paper "Infinitesimal Plane-based Pose Estimation"
% by Toby Collins and Adrien Bartoli, published in the International Journal of Computer Vision, September 2014.
% A copy of the author's pre-print version can be found here: http://isit.u-clermont1.fr/~ab/Publications/Collins_Bartoli_IJCV14.pdf
%
% Feel free to contact Toby (toby.collins@gmail.com) if you have any
% questions about the paper and IPPE.
%
% This package is free and covered by the BSD licence without any warranty. We hope you find this code useful and please cite our paper in your work:
% (c) Toby Collins 2015
%
%@article{
%year={2014},
%issn={0920-5691},
%journal={International Journal of Computer Vision},
%volume={109},
%number={3},
%doi={10.1007/s11263-014-0725-5},
%title={Infinitesimal Plane-Based Pose Estimation},
%url={http://dx.doi.org/10.1007/s11263-014-0725-5},
%publisher={Springer US},
%keywords={Plane; Pose; SfM; PnP; Homography},
%author={Collins, Toby and Bartoli, Adrien},
%pages={252-286},
%language={English}
%}

opts = parseOpts(opts);

measureTiming = opts.measureTiming;
refineSolutions = opts.withPoseRefinement;

U = double(U);
Q = double(Q);

%set defaults for hEstMethod:
if isempty(hEstMethod)
    hEstMethod = 'Harker';
else
    if ~(strcmp(hEstMethod,'Harker')||strcmp(hEstMethod,'DLT'))
        error('You must specify a valid homography estimation method in perspectiveIPPE. The options are Harker or DLT');
    end
end

%argument sanity check:
assert(size(U,1)==2|size(U,1)==3);
assert(size(U,2)==size(Q,2));
assert(size(Q,1) == 2);

%n is the number of points:
n = size(U,2);
modelDims = size(U,1);

Uuncentred = U;

if modelDims==2
   %we zero-center the model points:
    Pbar = [mean(U,2);0];
    U(1,:) = U(1,:)-Pbar(1);
    U(2,:) = U(2,:)-Pbar(2);
else
    %we rotate the model points onto the plane z=0 and zero center them:
    Pbar = mean(U,2);
    MCenter = eye(4);
    MCenter(1:3,end) = -Pbar;
    U_ = MCenter(1:3,:)*[U;ones(1,size(U,2))];
    %compute the rotation that rotates the model points onto the plane z=0:
    [modelRotation,sigs,~] = svd(U_*U_'); modelRotation=modelRotation';
    if det(modelRotation)<0
        modelRotation(3,:)  = -modelRotation(3,:); %this ensures modelRotation is a member of SO3
    end
    if (sigs(3,3)/sigs(2,2)>1e-5)
        error('IPPE requires the model points to be planar!');
    end

    modelRotation(4,4) = 1;
    Mcorrective = modelRotation*MCenter;
    U = Mcorrective(1:2,:)*[U;ones(1,size(U,2))];

end

if measureTiming
    tic;
end
%compute the model-to-image homography:
switch hEstMethod
    case 'DLT'
        H = homography2d([U;ones(1,n)],[Q;ones(1,n)]);
    case 'Harker'
        H = homographyHarker([U;ones(1,n)],[Q;ones(1,n)]);
end

%[ px, py, 1, 0, 0, 0, -px*qx, -py*qx]h = qx
%[ 0, 0, 0, px, py, 1, -px*qy, -py*qy]h = qy

%Compute the Jacobian J of the homography at (0,0):
H = H./H(3,3);
J(1,1) = H(1,1)-H(3,1)*H(1,3);
J(1,2) = H(1,2)-H(3,2)*H(1,3);
J(2,1) = H(2,1)-H(3,1)*H(2,3);
J(2,2) = H(2,2)-H(3,2)*H(2,3);

%Compute the two rotation solutions:
v = [H(1,3);H(2,3)];
[R1,R2] = IPPE_dec(v,J);

%Compute the two translation solutions:
t1_ = estT(R1,U,Q);
t2_ = estT(R2,U,Q);

if modelDims==2
    t1 = [R1,t1_]*[-Pbar;1];
    t2 = [R2,t2_]*[-Pbar;1];
else
    M1 = [R1,t1_];
    M1(4,4) = 1;
    M2 = [R2,t2_];
    M2(4,4) = 1;
    M1 = M1*Mcorrective;
    M2 = M2*Mcorrective;
    R1 = M1(1:3,1:3);
    R2 = M2(1:3,1:3);
    t1 = M1(1:3,end);
    t2 = M2(1:3,end);
end

[reprojErr1,reprojErr2] = computeReprojErrs(R1,R2,t1,t2,Uuncentred,Q);
%sort solutions:
if reprojErr1>reprojErr2
    [R1,R2,t1,t2,reprojErr1,reprojErr2] = swapSolutions(R1,R2,t1,t2,reprojErr1,reprojErr2);
end
IPPEPoses.R1 = R1;
IPPEPoses.t1 = t1;
IPPEPoses.R2 = R2;
IPPEPoses.t2 = t2;
IPPEPoses.reprojError1 = reprojErr1;
IPPEPoses.reprojError2 = reprojErr2;

%%with timing?
if measureTiming
    disp('Computing time to estimate homography for IPPE. Warning this slows things down, so turn it off when you are not benchmarking!');
    tic;
    %average timing over 100 runs:
    for i=1:100
        switch hEstMethod
            case 'DLT'
                H = homography2d([U;ones(1,n)],[Q;ones(1,n)]);
            case 'Harker'
                H = homographyHarker([U;ones(1,n)],[Q;ones(1,n)]);
        end
    end

    tm = toc;
    disp(['Time to estimate homography: ' num2str(1000*tm/100) 'ms']);

    disp('Computing time to solve IPPE given homography...');
    tic;
    %average timing over 100 runs:
    for i=1:100
        %Compute the Jacobian J of the homography at (0,0):
        H = H./H(3,3);
        J(1,1) = H(1,1)-H(3,1)*H(1,3);
        J(1,2) = H(1,2)-H(3,2)*H(1,3);
        J(2,1) = H(2,1)-H(3,1)*H(2,3);
        J(2,2) = H(2,2)-H(3,2)*H(2,3);
        %Compute the two rotation solutions:
        v = [H(1,3);H(2,3)];
        [R1,R2] = IPPE_dec(v,J);
        %Compute the two translation solutions:
        t1_ = estT(R1,U,Q);
        t2_ = estT(R2,U,Q);
        t1 = [R1,t1_]*[-Pbar;1];
        t2 = [R2,t2_]*[-Pbar;1];
    end
    tm = toc;
    disp(['Time to solve IPPE given homography: ' num2str(1000*tm/100) 'ms']);
end

if modelDims==2
    Uuncentred(3,:) = 0;
end

if refineSolutions
    poseEst= pnpLMRefinement(Uuncentred,Q,R1,t1,'MLE');
    R1 = poseEst.R;
    t1 = poseEst.t;

    poseEst= pnpLMRefinement(Uuncentred,Q,R2,t2,'MLE');
    R2 = poseEst.R;
    t2 = poseEst.t;

    [reprojErr1,reprojErr2] = computeReprojErrs(R1,R2,t1,t2,Uuncentred,Q);
    %Sort the solutions:
    if reprojErr1>reprojErr2
        [R1,R2,t1,t2,reprojErr1,reprojErr2] = swapSolutions(R1,R2,t1,t2,reprojErr1,reprojErr2);
    end
    refinedPoses.R1 = R1;
    refinedPoses.t1 = t1;
    refinedPoses.R2 = R2;
    refinedPoses.t2 = t2;
    refinedPoses.reprojError1 = reprojErr1;
    refinedPoses.reprojError2 = reprojErr2;

    %%with timing?
    if measureTiming
        disp('Computing time to refine IPPE solutions with Levenberg Marquardt...');
        tic;
        %average timing over 100 runs:
        for i=1:100
            pnpLMRefinement(Uuncentred,Q,IPPEPoses.R1,IPPEPoses.t1,'MLE');
            pnpLMRefinement(Uuncentred,Q,IPPEPoses.R2,IPPEPoses.t2,'MLE');
        end
        tm = toc;
        disp(['Time to refine IPPE solutions with Levenberg Marquardt: ' num2str(1000*tm/100) 'ms']);
    end

end

function [reprojErr1,reprojErr2] = computeReprojErrs(R1,R2,t1,t2,U,Q)
%computeReprojErrs: Computes the reprojection errors for the two solutions
%generated by IPPE.

%transform model points to camera coordinates and project them onto the
%image:
if size(U,1)==2
    PCam1 = R1(:,1:2)*U;
    PCam2 = R2(:,1:2)*U(1:2,:);
else
    PCam1 = R1*U;
    PCam2 = R2*U;
end

PCam1(1,:) = PCam1(1,:) + t1(1);
PCam1(2,:) = PCam1(2,:) + t1(2);
PCam1(3,:) = PCam1(3,:) + t1(3);

PCam2(1,:) = PCam2(1,:) + t2(1);
PCam2(2,:) = PCam2(2,:) + t2(2);
PCam2(3,:) = PCam2(3,:) + t2(3);

Qest_1 = PCam1./[PCam1(3,:);PCam1(3,:);PCam1(3,:)];
Qest_2 = PCam2./[PCam2(3,:);PCam2(3,:);PCam2(3,:)];

%compute reprojection errors:
reprojErr1 = norm(Qest_1(1:2,:)-Q);
reprojErr2 = norm(Qest_2(1:2,:)-Q);

function t = estT(R,psPlane,Q)
%Computes the least squares estimate of translation given the rotation solution.
if size(psPlane,1)==2
    psPlane(3,:) =0;
end
%qq = homoMult(Kinv,pp')';
Ps = R*psPlane;

numPts = size(psPlane,2);
Ax = zeros(numPts,3);
bx = zeros(numPts,1);

Ay = zeros(numPts,3);
by = zeros(numPts,1);

Ax(:,1) = 1;
Ax(:,3) = -Q(1,:);
bx(:) = Q(1,:).*Ps(3,:) -  Ps(1,:);

Ay(:,2) = 1;
Ay(:,3) = -Q(2,:);
by(:) = Q(2,:).*Ps(3,:) -  Ps(2,:);

A = [Ax;Ay];
b = [bx;by];

AtA = A'*A;
Atb = A'*b;

Ainv = IPPE_inv33(AtA);
t = Ainv*Atb;

function [R1,R2,t1,t2,reprojErr1,reprojErr2] = swapSolutions(R1_,R2_,t1_,t2_,reprojErr1_,reprojErr2_)
%swaps the solutions:
R1 = R2_;
t1 = t2_;
reprojErr1 = reprojErr2_;

R2 = R1_;
t2 = t1_;
reprojErr2 = reprojErr1_;

function Ainv = IPPE_inv33(A)
%computes the inverse of a 3x3 matrix, assuming it is full-rank.
a11 = A(1,1);
a12 = A(1,2);
a13 = A(1,3);

a21 = A(2,1);
a22 = A(2,2);
a23 = A(2,3);

a31 = A(3,1);
a32 = A(3,2);
a33 = A(3,3);

Ainv = [[ a22*a33 - a23*a32, a13*a32 - a12*a33, a12*a23 - a13*a22];
    [ a23*a31 - a21*a33, a11*a33 - a13*a31, a13*a21 - a11*a23];
    [ a21*a32 - a22*a31, a12*a31 - a11*a32, a11*a22 - a12*a21]];
Ainv = Ainv./(a11*a22*a33 - a11*a23*a32 - a12*a21*a33 + a12*a23*a31 + a13*a21*a32 - a13*a22*a31);

function opts = parseOpts(opts)
if ~isfield(opts,'measureTiming')
    opts.measureTiming = false;
end
if ~isfield(opts,'withPoseRefinement')
    opts.withPoseRefinement = true;
end


function [p1,p2,dp1,dp2] = autoGen_cartPoleKinematics(x,q,dx,dq,l,empty)
%AUTOGEN_CARTPOLEKINEMATICS
%    [P1,P2,DP1,DP2] = AUTOGEN_CARTPOLEKINEMATICS(X,Q,DX,DQ,L,EMPTY)

%    This function was generated by the Symbolic Math Toolbox version 6.3.
%    30-Nov-2015 20:13:07

p1 = [x;empty];
if nargout > 1
    t2 = cos(q);
    t3 = sin(q);
    p2 = [x+l.*t3;-l.*t2];
end
if nargout > 2
    dp1 = [dx;empty];
end
if nargout > 3
    dp2 = [dx+dq.*l.*t2;dq.*l.*t3];
end


function tapas_hgf_ar1_binary_mab_plotTraj(r)
% Plots the estimated or generated trajectories for the binary HGF perceptual model for multi-armed
% bandit situations.
%
% Usage example:  est = tapas_fitModel(responses, inputs); tapas_hgf_binary_plotTraj(est);
%
% --------------------------------------------------------------------------------------------------
% Copyright (C) 2013 Christoph Mathys, TNU, UZH & ETHZ
%
% This file is part of the HGF toolbox, which is released under the terms of the GNU General Public
% Licence (GPL), version 3. You can redistribute it and/or modify it under the terms of the GPL
% (either version 3 or, at your option, any later version). For further details, see the file
% COPYING or <http://www.gnu.org/licenses/>.

% Optional plotting of standard deviations (true or false)
plotsd = true;

% Optional plotting of responses (true or false)
ploty = true;

% Set up display
scrsz = get(0,'screenSize');
outerpos = [0.2*scrsz(3),0.2*scrsz(4),0.8*scrsz(3),0.8*scrsz(4)];
figure(...
    'OuterPosition', outerpos,...
    'Name', 'HGF trajectories');

% Set up colors
colors = [1 0 0; 0.67 0 1; 0 0.67 1; 0.67 1 0];

% Number of bandits
b = r.c_prc.n_bandits;

% Number of trials
n = size(r.u,1);

% Time axis
t = ones(1,n);

ts = cumsum(t);
ts = [0, ts];

% Number of levels
try
    l = r.c_prc.n_levels;
catch
    l = (length(r.p_prc.p)+1)/5;
end

% Upper levels
for j = 1:l-2

    % Subplots
    subplot(l,1,j);

    if plotsd == true
        upperprior = r.p_prc.mu_0(l-j+1) +sqrt(r.p_prc.sa_0(l-j+1));
        lowerprior = r.p_prc.mu_0(l-j+1) -sqrt(r.p_prc.sa_0(l-j+1));
        upper = [upperprior; r.traj.mu(:,l-j+1,1)+sqrt(r.traj.sa(:,l-j+1,1))];
        lower = [lowerprior; r.traj.mu(:,l-j+1,1)-sqrt(r.traj.sa(:,l-j+1,1))];

        plot(0, upperprior, 'ob', 'LineWidth', 1);
        hold all;
        plot(0, lowerprior, 'ob', 'LineWidth', 1);
        fill([ts, fliplr(ts)], [(upper)', fliplr((lower)')], ...
             'b', 'EdgeAlpha', 0, 'FaceAlpha', 0.15);
    end
    plot(ts, [r.p_prc.mu_0(l-j+1); r.traj.mu(:,l-j+1,1)], 'b', 'LineWidth', 2);
    hold all;
    plot(0, r.p_prc.mu_0(l-j+1), 'ob', 'LineWidth', 2); % prior
    xlim([0 ts(end)]);
    title(['Posterior expectation of x_' num2str(l-j+1)], 'FontWeight', 'bold');
    ylabel(['\mu_', num2str(l-j+1)]);
end

% Second level
subplot(l,1,l-1)

if plotsd == true
    for j=1:b
        upperprior = r.p_prc.mu_0(2) +sqrt(r.p_prc.sa_0(2));
        lowerprior = r.p_prc.mu_0(2) -sqrt(r.p_prc.sa_0(2));
        upper = [upperprior; r.traj.mu(:,2,j)+sqrt(r.traj.sa(:,2,j))];
        lower = [lowerprior; r.traj.mu(:,2,j)-sqrt(r.traj.sa(:,2,j))];

        plot(0, upperprior, 'o', 'Color', colors(j,:), 'LineWidth', 1);
        hold all;
        plot(0, lowerprior, 'o', 'Color', colors(j,:), 'LineWidth', 1);
        fill([ts, fliplr(ts)], [(upper)', fliplr((lower)')], ...
             colors(j,:), 'EdgeAlpha', 0, 'FaceAlpha', 0.15);
    end
end
for j=1:b
    plot(ts, [r.p_prc.mu_0(1); r.traj.mu(:,2,j)], 'Color', colors(j,:), 'LineWidth', 2);
    hold all;
    plot(0, r.p_prc.mu_0(1), 'o', 'Color', colors(j,:), 'LineWidth', 2); % prior
end
xlim([0 ts(end)]);
title('Posterior expectations of x_2', 'FontWeight', 'bold');
ylabel('\mu_2');

% Input level
subplot(l,1,l);

for j=1:b
    plot(ts, [tapas_sgm(r.p_prc.mu_0(2), 1); tapas_sgm(r.traj.mu(:,2,j), 1)], 'Color', colors(j,:), 'LineWidth', 2);
    hold all;
    plot(0, tapas_sgm(r.p_prc.mu_0(2), 1), 'o', 'Color', colors(j,:), 'LineWidth', 2); % prior
end
plot(ts(2:end), r.u(:,1), '.', 'Color', [0 0 0]); % inputs
plot(ts(2:end), r.traj.wt(:,1), 'k') % implied learning rate
if (ploty == true) && ~isempty(find(strcmp(fieldnames(r),'y'))) && ~isempty(r.y)
    y = r.y(:,1);
    if ~isempty(find(strcmp(fieldnames(r),'irr')))
        y(r.irr) = NaN; % weed out irregular responses
        plot(ts(r.irr)+1,  1.08.*ones([1 length(r.irr)]), 'x', 'Color', [1 0.7 0], 'Markersize', 11, 'LineWidth', 2); % irregular responses
    end
    for j=1:b
        plot(find(y==j), 1.08*ones([1 length(find(y==j))]), '.', 'Color', colors(j,:)); % responses
    end
    title(['Response y, input u (black dots), learning rate (fine black line), and posterior expectation of reward s(\mu_2) ', ...
           '(colour coded) for \phi=', num2str(r.p_prc.phi(2:end)), ', m=', num2str(r.p_prc.m(2:end)), ', \kappa=', ...
           num2str(r.p_prc.ka(2:end)), ', \omega=', num2str(r.p_prc.om(2:end))], ...
      'FontWeight', 'bold');
    ylabel('y, u, s(\mu_2)');
    axis([0 ts(end) -0.15 1.15]);
else
    title(['Input u (black dots), learning rate (fine black line), and posterior expectation of input s(\mu_2) ', ...
           '(red) for \rho=', num2str(r.p_prc.rho(2:end)), ', \kappa=', ...
           num2str(r.p_prc.ka(2:end)), ', \omega=', num2str(r.p_prc.om(2:end))], ...
      'FontWeight', 'bold');
    ylabel('u, s(\mu_2)');
    axis([0 ts(end) -0.1 1.1]);
end
plot(ts(2:end), 0.5, 'k');
xlabel('Trial number');
hold off;


function [n_signif,index_signif] = fdr(pvals,qlevel,method,adjustment_method,adjustment_args)
% Determine significance based on a false discovery rate (FDR) method
%
%  [n_signif,index_signif] = ...
%     fdr(pvals,qlevel,method,adjustment_method,adjustment_args)
%
% This is the main function designed for general usage for determining
% significance based on the FDR approach. The specific method used depends
% on the calling parameters.
%
% Arguments:
%
%   pvals:  a vector of pvals on which to conduct the multiple testing
%
%   qlevel: the proportion of false positives desired
%
%   method: method for performing the testing.  'original' follows
%   Benjamini & Hochberg (1995); 'general' is much more
%   conservative, requiring no assumptions on the p-values (see
%   Benjamini & Yekutieli (2001)).  We recommend using 'original',
%   and if desired, using 'adjustment.method="mean" ' to increase
%   power.
%
%   adjustment_method: method for increasing the power of the
%   procedure by estimating the proportion of alternative p-values,
%   one of 'mean', the modified Storey estimator that we suggest in
%   Ventura et al. (2004), 'storey', the method of Storey (2002),
%   or 'two-stage', the iterative approach of Benjamini et
%   al. (2001)
%
%   adjustment_args: arguments to adjustment.method; see prop_alt()
%   for description, but note that for 'two-stage', qlevel and
%   fdr_method are taken from the qlevel and method arguments to
%   fdr()
%
% Value:
%
%   index_signif:  a vector of the indices of the significant tests
%   or NA if no significant tests
%
%   n_signif: number of significant tests
%
% Examples:
%
%   signif <- fdr(pvals,method='original',adjustment.method='mean')
%
%  References:
%             Ventura, V., C.J. Paciorek, and J.S. Risbey.  2004.
%             Controlling the proportion of falsely-rejected
%             hypotheses when conducting multiple tests with
%             climatological data.  Journal of Climate, in press.
%             Also Carnegie Mellon University, Department of
%             Statistics technical report 775
%             (www.stat.cmu.edu/tr/tr775/tr775.html).
%
%             Benjamini, Y, and Y. Hochberg. 1995. Controlling the
%             false discovery rate: a practical and powerful
%             approach to multiple testing.  JRSSB 57:289-300.
%
%             Benjamini, Y. and D. Yekutieli.  2001.  The control
%             of the false discovery rate in multiple testing under
%             dependency. Annals of Statistics 29:1165-1188.
%
%             Benjamini, Y., A. Krieger, and D. Yekutieli.  2001.
%             Two staged linear step up FDR controlling procedure.
%             Technical Report, Department of Statistics and
%             Operations Research, Tel Aviv University.  URL:
%             http://www.math.tau.ac.il/~ybenja/Papers.html
%
%             Storey, J. 2002.  A direct approach to false discovery rates.  JRSSB 64: 479-498.
%
% Attribution:
%
%    Version 0.1.3;  May 10, 2004
%
%    License: GPL version 2 or later
%
% This program is free software; you can redistribute it and/or
% modify it under the terms of the GNU General Public License
% as published by the Free Software Foundation; either version 2
% of the License, or (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
%
%    Author: Chris Paciorek - please contact the author with bug
%    reports: paciorek AT alumni.cmu.edu
%

if exist('isfield') % matlab
    endline='';
else % octave
    endline='\n';
end

n = length(pvals);

a = 0;   % initialize proportion of alternative hypotheses
if nargin<2
    error(['pvals and qlevel are required arguments',endline]);
end
if nargin==2
    method='original';
end
if nargin>=4 && ~isempty(adjustment_method)
    if strcmp(adjustment_method,'two-stage')  % set things up for the "two-stage" estimator
        qlevel = qlevel/(1+qlevel);   % see Benjamini et al. (2001) for proof that this controls the FDR at level qlevel
        adjustment_args.qlevel = qlevel;
        adjustment_args.fdr_method = method;
        %fprintf(['Adjusting cutoff using two-stage approach with fdr ', ...
        %    'method %s and qlevel %.3f\n'],adjustment_args.fdr_method,...
        %    adjustment_args.qlevel);
    end
    if strcmp(adjustment_method,'mean') & nargin==4  % default arguments for "mean" method of Ventura et al. (2004)
        adjustment_args.edf_lower=0.8;
        adjustment_args.num_steps=20;
        %fprintf(['Adjusting cutoff using mean approach with edf_lower=0.8 ', ...
        %    'and num_steps=20\n']);
    end
    if strcmp(adjustment_method,'storey') & nargin==4
        error(['adjustment_args.edf_quantile not specified for Storey ', ...
            'adjustment',endline]);
    end
    a = prop_alt(pvals,adjustment_method,adjustment_args);
end
if a==1    % all hypotheses are estimated to be alternatives
    index_signif=1:n;
    n_signif=n;
else      % adjust for estimate of a; default is 0
    qlevel = qlevel/(1-a);
    [n_signif,index_signif]=fdr_master(pvals,qlevel,method);
end

function [n_signif,index_signif] = fdr_master(pvals,qlevel,method)
%
% Description:
%
%    This is an internal function that performs various versions of
%    the FDR procedure, but without the modification described in
%    section 4 of Ventura et al. (2004)
%
% Arguments:
%
%   pvals (required):  a vector of pvals on which to conduct the multiple testing
%
%   qlevel: the proportion of false positives desired
%
%   method: one of 'original', the original method of Benjamini &
%   Hochberg (1995), or 'general', the method of Benjamini &
%   Yekutieli (2001), which requires no assumptions about the
%   p-values, but which is much more conservative.  We recommend
%   'original' for climatological data, and suspect it works well
%   generally for spatial data.
%
% Value:
%
%   index_signif:  a vector of the indices of the significant tests
%   or NA if no significant tests
%
%   n_signif: number of significant tests
%
%
if nargin==2
    method='original';
end
n = length(pvals);
if strcmp(method,'general')
    qlevel = qlevel/sum(1./(1:n));  % This requires fewer assumptions but is much more conservative
else
    if(~strcmp(method,'original'))
        error('No method of type: %s\n',method);
    end
end
[n_signif,index_signif] = fdr_basic(pvals,qlevel);

function [n_signif,index_signif] = fdr_basic(pvals,qlevel)
%
% Description:
%
%    This is an internal function that performs the basic FDR of Benjamini & Hochberg (1995).
%
% Arguments:
%
%   pvals:  a vector of pvals on which to conduct the multiple testing
%
%   qlevel: the proportion of false positives desired
%
% Value:
%
%   index_signif:  a vector of the indices of the significant tests
%   or NA if no significant tests
%
%   n_signif: number of significant tests
%
n = length(pvals);
if n>1 & size(pvals,2)==1
    pvals=pvals';
end
[sorted_pvals,sort_index] = sort(pvals);
indices = (1:n).*(sorted_pvals<=qlevel*(1:n)/n);
n_signif = max(indices);
if(n_signif)
    indices = 1:n_signif;
    index_signif=sort(sort_index(indices));
else
    index_signif=nan;
end

function [a] = storey(edf_quantile,pvals)
%
% Description:
%
%    This is an internal function that calculates the basic Storey
%    (2002) estimator of a, the proportion of alternative
%    hypotheses.
%
% Arguments:
%
%   edf_quantile (required):  the quantile of the empirical
%   distribution function at which to estimate a
%
%   pvals (required):  a vector of pvals on which to estimate a
%
% Value:
%
%   a: estimate of a, the number of alternative hypotheses
%
%
if exist('isfield') % matlab
    endline='';
else % octave
    endline='\n';
end
if edf_quantile >= 1 | edf_quantile <=0
    error(['edf_quantile should be between 0 and 1',endline]);
end
a = (mean(pvals<=edf_quantile)-edf_quantile)/(1-edf_quantile);
if(a<0)
    a=0;
end

function [a] = prop_alt(pvals,adjustment_method,adjustment_args)
%
% Description:
%
%    This is an internal function that calculates an estimate of a,
%    the proportion of alternative hypotheses, using one of several
%    methods.
%
% Arguments:
%
%   pvals (required):  a vector of pvals from which to estimate a
%
%   adjustment_method: method for  estimating the proportion of
%   alternative p-values, one of "mean", the modified Storey
%   estimator suggested in Ventura et al. (2004); "storey", the
%   method of Storey (2002); or "two-stage", the iterative approach
%   of Benjamini et al. (2001)
%
%   adjustment_args: arguments to adjustment.method;
%
%      for 'mean', specify edf_lower, the smallest quantile at
%      which to estimate a, and num_steps, the number of quantiles
%      to use - the approach uses the average of the Storey (2002)
%      estimator for the num.steps quantiles starting at edf.lower
%      and finishing just less than 1
%
%      for 'storey', specify edf_quantile, the quantile at which to
%      calculate the estimator
%
%      for 'two-stage', the method uses a standard FDR approach to
%      estimate which p-values are significant; this number is the
%      estimate of a; therefore the method requires specification
%      of qlevel, the proportion of false positives and fdr.method
%      ('original' or 'general'), the FDR method to be used.  We do
%      not recommend 'general' as this is very conservative and
%      will underestimate a.

%
% Value:
%
%   a:  estimate of a, the number of alternative hypotheses
%
% Examples:
%
%   a = prop_alt(pvals,'mean')
%
if exist('isfield') % matlab
    endline='';
else % octave
    endline='\n';
end

stop=0;
n = length(pvals);
if strcmp(adjustment_method,'two-stage')
    if ~exist('adjustment_args')
        stop=1;
    else
        if exist('isfield')  % check if in Matlab
            if ~(isfield(adjustment_args,'qlevel') & ...
                    isfield(adjustment_args,'fdr_method'))
                stop=1;
            end
        else  % in octave
            if ~(struct_contains(adjustment_args,'qlevel') & ...
                    struct_contains(adjustment_args,'fdr_method'))
                stop=1;
            end
        end
    end
    if stop
        error(['adjustment_args.qlevel or adjustment_args.fdr_method ', ...
            'not specified.  Two-stage estimation of the number ',...
            'of alternative hypotheses requires specification of ', ...
            'the  FDR threshold and FDR method',endline]);
    end
    [n_signif,index_signif]=fdr_master(pvals,adjustment_args.qlevel,...
        adjustment_args.fdr_method);
    a=n_signif/n;
end
if strcmp(adjustment_method,'storey')
    if ~(exist('adjustment_args'))
        stop=1;
    else
        if exist('isfield')  % check if in Matlab
            if ~(isfield(adjustment_args,'edf_quantile'))
                stop=1;
            end
        else  % octave
            if ~(struct_contains(adjustment_args,'edf_quantile'))
                stop=1;
            end
        end
    end
    if stop
        error(['adjustment_args.edf_quantile not specified. Using method of Storey for estimating  the number of alternative hypotheses requires specification of the argument of the p-value EDF at which to do the estimation (a number close to one is recommended)',endline]);
    end
    a=storey(adjustment_args.edf_quantile,pvals);
end

if strcmp(adjustment_method,'mean')
    if ~(exist('adjustment_args'))
        stop=1;
    else
        if exist('isfield')  % check if in Matlab
            if ~(isfield(adjustment_args,'edf_lower') & ...
                    isfield(adjustment_args,'num_steps'))
                stop=1;
            end
        else % octave
            if ~(struct_contains(adjustment_args,'edf_lower') & ...
                    struct_contains(adjustment_args,'num_steps'))
                stop=1;
            end
        end
    end
    if stop
        error(['adjustment_args.edf_lower or adjustment_args.num_steps ' ...
            'is not specified. Using the method of Ventura et al. (2004) for estimating  the number of alternative hypotheses requires specification of the lowest quantile of the p-value EDF at which to do the estimation (a number close to one is recommended) and the number of steps between edf_lower and 1, starting at edf_lower, at which to do the estimation',endline])
    end
    if adjustment_args.edf_lower >=1 | adjustment_args.edf_lower<=0
        error(['adjustment_args.edf_lower must be between 0 and 1',endline]);
    end
    if adjustment_args.num_steps<1 | ~(rem(adjustment_args.num_steps,1)==0)
        error(['adjustment_args.num_steps must be an integer greater than 0',endline]);
    end
    stepsize = (1-adjustment_args.edf_lower)/adjustment_args.num_steps;
    edf_quantile=adjustment_args.edf_lower;
    a_vec=storey(edf_quantile,pvals);
    for step=1:adjustment_args.num_steps
        a_vec(step)=storey(edf_quantile,pvals);
        edf_quantile=edf_quantile+stepsize;
    end
    a=mean(a_vec);
end


function [] = aps_wrf_SAR()
% [] = aps_wrf_SAR()
% Script to load wrf processed data, on 37 pressure levels. Load data for area of interest,
% and calculate delays maps for each specified date.
% The DEM file inputed should have an asociated
% ".rsc" file, with the same filename as the DEM. The ".rsc" files should
% contain a WIDTH, LENGTH, X_FIRST, Y_FIRST, X_STEP, Y_STEP and optional a
% FORMAT string. The meris data is assumed to be structured in date folders.
% The batchfile contains the full path to the meris files in these folders.
% Note that the first line of the batchfile should read "files".
%
%
% INPUTS:
% demfile               Full path to the DEM file. The DEM needs to me in meters.
% xlims                 Limits in the x-direction, either in degrees
% ylims                 Limits in the y-direction, either in degrees
% demnull               The value for no DEM data, default is -32768.
% smpres                The output resolution, either in degrees
%                       Units needs to be consistend with xlims and ylims.
%
% OUTPUTS:
% It will give the computed ZENITH dry and wet delay map in cm for the selected region.
%
%     Copyright (C) 2015  Bekaert David - University of Leeds
%     Email: eedpsb@leeds.ac.uk or davidbekaert.com
%
%     This program is free software; you can redistribute it and/or modify
%     it under the terms of the GNU General Public License as published by
%     the Free Software Foundation; either version 2 of the License, or
%     (at your option) any later version.
%
%     This program is distributed in the hope that it will be useful,
%     but WITHOUT ANY WARRANTY; without even the implied warranty of
%     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%     GNU General Public License for more details.
%
%     You should have received a copy of the GNU General Public License along
%     with this program; if not, write to the Free Software Foundation, Inc.,
%     51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
%
% Richard Walters/ Bekaert David - Oxford / Leeds University 2012-2013
% modifications:
% 19/11/2013    DB      Change the computation of the saturated vapor
%                       pressure over ice and water
% 19/11/2013    DB      Change the computation of the hydrostatic delay to
%                       be directly computed from pressure and temperature
% 20/11/2013    DB      Check if wrf data exist to avoid crashing
% 21/11/2013    DB      Interpolate H,Temp, Hum spatially to cope with
%                       coarser land mask compared to the DEM
% 02/01/2014	DB      fix for NaN value over sea
% 04/08/2014    DB      Redefine meris_lat(lon)_range to region_lat(lon)_range
% 24/02/2015    DB      Save support information for WRF
% 27/09/2015    DB      Separate the DEM loading, ask for number of domains.
% 28/09/2015    DB      Also save the delay computation for a grid or z to include Pirate compartibility.
% 28/09/2015    DB      Include multi-core from matlab
% 07/07/2016    DB      Redefine hydrostatic delay to be based on surface pressure.

save_complete=0;        % save support information when 0

%% Constants
% Parameters from Jolivet et al 2011, GRL
Rd = 287.05;            % [j/kg/K] specific gas constants dry air
Rv = 461.495;           % [j/kg/K] water vapour
k1 = 0.776;             % [K/Pa]
k2 = 0.716;             % [K/Pa]
k3 = 3.75e3;            % [K^2/Pa]

%% Defaults
zref = 15000;       % zref for integration- tropopause
zincr = 15;         % z increment spacing for vertical interpolation before integration
vertres = 5;        % vertical resolution for comparing dem to vert profiles of delay

%% get the number of domains
repeat=1;
while repeat==1
    action_flag= str2num(input('How many domains do you have including the parent? ','s'));
    if isnumeric(action_flag)
        n_domains = action_flag;
        repeat=0;
    end
end

%% getting the variables from the parms_aps file
stamps_processed = getparm_aps('stamps_processed',1);
ll_matfile = getparm_aps('ll_matfile',1);
wrf_datapath = getparm_aps('wrf_datapath',1);
datestructure = 'yyyymmdd';                               % assumed date structure for ERA-I
UTC_sat =  getparm_aps('UTC_sat',1);

% loading the data
if strcmp(stamps_processed,'y')
   ps = load(ll_matfile);
   dates = ps.day;
   load psver
   fprintf('Stamps processed structure \n')
else
    psver = 2;
    ifgday_matfile = getparm_aps('ifgday_matfile',1);
    ifgs_dates = load(ifgday_matfile);
    ifgs_dates = ifgs_dates.ifgday;
    dates = reshape(ifgs_dates,[],1);
    dates = unique(dates);
    dates = datenum(num2str(dates),'yyyymmdd');

end

% getting the dates
n_dates = length(dates);

%% Compute and resample DEM
[dem,xmin,xmax,ymin,ymax,smpres,nncols,nnrows] = get_DEM;

% the region which is cropped from the ERA data and used to make the interpolation.
% Should be  larger than the region to which the delay is computed
lonmin = floor(xmin)-1;
lonmax= ceil(xmax)+1;
latmin = floor(ymin)-1;
latmax = ceil(ymax)+1;

% setting the maximum height of the DEM to limit the range at whcih ERA-!
% needs to be interpolated to
maxdem = ceil(max(max(dem))/100)*100+50; % max height of dem in metres

fprintf(['Interpolate to a maximum dem height of ', num2str(maxdem) ,' m\n'])

%% performing the calucluation for each  date

for d = 1:n_dates
    date = datestr(dates(d),datestructure);
    date_str_file = datestr(dates(d),'yyyy-mm-dd');

    % the netcdf files
    infile = [wrf_datapath filesep date filesep 'wrfplev_d0' num2str(n_domains) '_' date_str_file '_' UTC_sat ':00'];
    infile2 = [wrf_datapath filesep date filesep 'wrfout_d0' num2str(n_domains) '_' date_str_file '_' UTC_sat ':00'];

    if exist(infile,'file')==2
        % the output file names
        outfile = [wrf_datapath filesep date filesep date '_ZWD.xyz'];
        hydroutfile = [wrf_datapath filesep date filesep date '_ZHD.xyz'];
        outfile2 = [wrf_datapath filesep date filesep date '_ZWD.z'];
        hydroutfile2 = [wrf_datapath filesep date filesep date '_ZHD.z'];

        % getting the information from the output file
        ncid = netcdf.open(infile,'NC_NOWRITE');
        [numdims,numvars,numglobalatts,unlimdimid] = netcdf.inq(ncid);
        [dimname, dimlen] = netcdf.inqDim(ncid,0); % dimensions
        % 3 = pressure*37
        % 6 = temperature

        % variables at each node 20 (0-19)
        [varname, vartype, dimids, natts] = netcdf.inqVar(ncid,0);

        % load variables
        Temp = netcdf.getVar(ncid,6);       % temp in K
        Hum = netcdf.getVar(ncid,7);        % relative humidity in percent
        H = netcdf.getVar(ncid,8);          % Geopotential Height in m
        Plevs = netcdf.getVar(ncid,3);      % 37 pressure levels in Pa
        netcdf.close(ncid)

        % flip the N-S axis
        Temp = flipdim(Temp,2);
        Temp = permute(Temp,[2,1,3]);
        Hum = flipdim(Hum,2);               % flip the N-S axis
        Hum = permute(Hum,[2,1,3]);
        H = flipdim(H,2);                   % flip the N-S axis
        H = permute(H,[2,1,3]);

        % replacing the non-usefull values by NaN
        Temp(Temp==-999)=NaN;
        Hum(Hum==-999)=NaN;
        H(H==-999)=NaN;

        % getting the longitude and latitude coordinates
        numlat = size(Temp,1);
        numlon = size(Temp,2);
        ncid = netcdf.open(infile2,'NC_NOWRITE');
        lats = netcdf.getVar(ncid,1);
        latgrid = repmat(lats,[1,1,37]);
        latgrid = permute(latgrid,[2,1,3]);
        latgrid = flipdim(latgrid,1);
        lons = netcdf.getVar(ncid,2);
        longrid = repmat(lons,[1,1,37]);
        longrid = permute(longrid,[2,1,3]);
        longrid = flipdim(longrid,1);
        % close netcdf file
        netcdf.close(ncid)

        % Load pressure levels and replicate vector to make 3d matrix same size as variables
        Plevs = Plevs./100; % convert into hPa (or millibars) 1000 hPa at surface
        Pressure = repmat(Plevs,[1,numlat,numlon]);
        Pressure = permute(Pressure,[2,3,1]);

        % Get list of points to look at in analysis
        [xx,yy] = meshgrid(1:numlon,1:numlat);
        latlist = reshape(latgrid(:,:,1),[],1);
        lonlist = reshape(longrid(:,:,1),[],1);
        xlist = reshape(xx,[],1);
        ylist = reshape(yy,[],1);

        % When the landmask is coarser it can introduce extreme values at
        % sharp transitions. To cope with this I perform an interpoaltion
        % over the land points, such the land mask is removed.
        % interpolate the land
        n_grid_points = size(Pressure,1)*size(Pressure,2);
        ix_vert =find(sum(sum(isnan(Temp)))==n_grid_points);
        if isempty(ix_vert)
            ix_vert = size(Temp,3);
        end
        longrid_vector = reshape(longrid(:,:,1),[],1);
        latgrid_vector= reshape(latgrid(:,:,1),[],1);

        % check for land mask only those layers that have land within them.
        parfor kkk=1:ix_vert(1)-1
             ix = find(isnan(Temp(:,:,kkk))==1);
             if ~isempty(ix)
                 Temp_temp = reshape(Temp(:,:,kkk),[],1);
                 Hum_temp = reshape(Hum(:,:,kkk),[],1);
                 H_temp = reshape(H(:,:,kkk),[],1);

                 longrid_vector_temp = longrid_vector;
                 latgrid_vector_temp = latgrid_vector;

                 % the land points that need to be interpolated
                 longrid_vector_needed = longrid_vector(ix);
                 latgrid_vector_needed = latgrid_vector(ix);

                 % removing the land poitns with currently NaN values
                 Temp_temp(ix)=[];
                 Hum_temp(ix)=[];
                 H_temp(ix)=[];

                 longrid_vector_temp(ix)=[];
                 latgrid_vector_temp(ix)=[];

                 % do the interpolation for the landpoints
                 Temp_vector_needed = griddata(double(longrid_vector_temp),double(latgrid_vector_temp),double(Temp_temp),double(longrid_vector_needed),double(latgrid_vector_needed),'linear');
                 Hum_vector_needed = griddata(double(longrid_vector_temp),double(latgrid_vector_temp),double(Hum_temp),double(longrid_vector_needed),double(latgrid_vector_needed),'linear');
                 H_vector_needed = griddata(double(longrid_vector_temp),double(latgrid_vector_temp),double(H_temp),double(longrid_vector_needed),double(latgrid_vector_needed),'linear');
                 Temp_new = Temp(:,:,kkk);
                 Hum_new = Hum(:,:,kkk);
                 H_new = H(:,:,kkk);

                 % replacing the values
                 Temp_new(ix)=Temp_vector_needed;
                 Hum_new(ix)=Hum_vector_needed;
                 H_new(ix)=H_vector_needed;

                 % storing them back in the array
                 Temp(:,:,kkk) = Temp_new;
                 Hum(:,:,kkk) = Hum_new;
                 H(:,:,kkk) =  H_new;
             end
        end
        clear Temp_new Hum_new H_new ix Temp_temp Hum_temp H_temp longrid_vector_temp latgrid_vector_temp

        % select the grid points within the area of interest
        % the projection of the grid is not uniform for the longitude as the center is defined for a specific longitude coordinate
        latlist_temp = sort(abs(diff(latlist)));
        lat_res = abs(latlist_temp(ceil(length(latlist_temp)*0.95)))*1.5;

        lonlist_temp = diff(lonlist);
        lon_res = max(abs(lonlist_temp))*1.5;

        % make sure the grid selected is covered and can be interpolated by WRF. Add a resolution cell to make sure it is the case.
        ix = find(ymin-lat_res<= latlist & latlist <= ymax+lat_res  & xmin-lon_res<= lonlist & lonlist<= xmax+lon_res) ;
        xlist = xlist(ix);
        ylist = ylist(ix);
        latlist = latlist(ix);
        lonlist = lonlist(ix);
        latlist=double(latlist);
        lonlist=double(lonlist);

        numy = length(unique(latlist));
        numx = length(unique(lonlist));
        ulatlist = unique(latlist);
        ulonlist = unique(lonlist);
        uxlist = unique(xlist);
        uylist = unique(ylist);

        if save_complete==0
            % saving the information for support plotting
            wrf.wrf_lonlat = [lonlist latlist];
            wrf.region = [[xmin xmin xmax xmax xmin]' [ymin ymax ymax ymin ymin]'];
            deminfo.xmin = xmin;
            deminfo.xmax = xmax;
            deminfo.ymax = ymax;
            deminfo.ymin = ymin;
            deminfo.dem = dem;
            wrf.deminfo = deminfo;
            clear deminfo
            % checking if the file already exist. Yes append otherwiuze create it
            if exist('tca_support.mat','file')==2
              save('tca_support.mat','-append','wrf')
            else
              save('tca_support.mat','wrf')
            end
            save_complete=1;
        end

        % Could not find the wrf used equation as they appear to be mixed
        % with latent heat etc. Istead I used the equations used in ERA-I
        % (see IFS documentation part 2: Data assimilation (CY25R1)). Calculate
        % saturated water vapour pressure (svp) for water (svpw) using Buck 1881
        % and for ice (swpi) from Alduchow and Eskridge (1996) euation AERKi
        svpw = 6.1121.*exp((17.502.*(Temp-273.16))./(240.97+Temp-273.16));
        svpi = 6.1121.*exp((22.587.*(Temp-273.16))./(273.86+Temp-273.16));
        tempbound1 = 273.16; %0
        tempbound2 = 250.16; %-23

        svp = svpw;
        % Faster expression
        wgt = (Temp - tempbound2)/(tempbound1 - tempbound2);
        svp = svpi+ (svpw - svpi).*wgt.^2; % not sure why have wgt^2, Romain does it in his pyaps code.
        ix_bound1 = find(Temp > tempbound1);
        svp(ix_bound1) = svpw(ix_bound1);
        ix_bound2 = find(Temp < tempbound2);
        svp(ix_bound2) = svpi(ix_bound2);

        e = Hum./100.*svp;
        % NB, P = Pressure in hydrostatic equation, NOT dry partial pressure as
        % stated (incorrectly) in Jolivet et al. 2011.
        P = Pressure;

        %% Convert Geopotential to Geopotential Height and then to Geometric Height
        g0 = 9.80665;

        % map of g with latitude
        g = 9.80616.*(1 - 0.002637.*cosd(2.*latgrid) + 0.0000059.*(cosd(2.*latgrid)).^2);
        % map of Re with latitude
        Rmax = 6378137;
        Rmin = 6356752;
        Re = sqrt(1./(((cosd(latgrid).^2)./Rmax^2) + ((sind(latgrid).^2)./Rmin^2)));

        %% Calculate Geometric Height, Z
        Z = (H.*Re)./(g/g0.*Re - H);

        % Find middle of scene to work out glocal and Rlocal for use later
        midx = round(mean(uxlist));
        midy = round(mean(uylist));
        glocal = g(midy,midx,1);
        Rlocal = Re(midy,midx,1);

        %% Find middle of scene to work out glocal and Rlocal for use later
        cdslices = maxdem/vertres +1;
        cdstack = zeros(size(lonlist,1),cdslices);
        cdstack_dry = zeros(size(lonlist,1),cdslices);
        cdstack_wet = zeros(size(lonlist,1),cdslices);

        XI=(0:zincr:zref)';
        gh = glocal.*(Rlocal./(Rlocal+XI)).^2; %gravity with height for XI height range

        % Interpolate Temp P and e from 0:20:15000 m
        % then integrate using trapz to estimate delay as function of height
        parfor m = 1:size(lonlist,1);
            xn = xlist(m);
            yn = ylist(m);

            %interpolate at res zincr before doing integration
            X=double(squeeze(Z(yn,xn,:)));
            YP=double(squeeze(P(yn,xn,:)));
            Ye=double(squeeze(e(yn,xn,:)));
            YT=double(squeeze(Temp(yn,xn,:)));

            YT(isnan(Ye))=[];
            YP(isnan(Ye))=[];
            X(isnan(Ye))=[];
            Ye(isnan(Ye))=[];

            YeI = interp1(X,Ye,XI,'spline')*100; %change from hPa to Pa
            YPI = interp1(X,YP,XI,'spline')*100; %change from hPa to Pa
            YTI = interp1(X,YT,XI,'spline');

            tmp1 = ((k2-(Rd*k1/Rv)).*YeI./YTI + k3.*YeI./(YTI.^2));
            Lw = (10^-6).*-1*flipud(cumtrapz(flipud(XI),flipud(tmp1)));
            % L is zenith delay one way in metres
            tmp2 = k1.*YPI./YTI;
            %Ld = (10^-6).*-1*flipud(cumtrapz(flipud(XI),flipud(tmp2)));             % This is using P/T expression (Hanssen, 2001)
            gm = glocal;
            Ld = (10^-6).*((k1*Rd/gm).*(YPI - YPI(zref/zincr +1)));                 % This is P0 expression (Hanssen, 2001)

            % Interpolate important part (i.e. total delay at elevations
            % less than maxdem) at high res i.e. vertres, and put in cdstack.
            cdI=(0:vertres:maxdem)';
            LdI=interp1(XI,Ld,cdI,'spline');
            LwI=interp1(XI,Lw,cdI,'spline');

            %cdstack(j,i,:) = LsI;
            cdstack_dry(m,:) = LdI;
            cdstack_wet(m,:) = LwI;
        end

        % Interpolate each cdstack layer onto a grid given by the DEM extents
        % in UTM m.
        xsmpres = (xmax-xmin)/nncols;
        ysmpres = (ymax-ymin)/nnrows;
        [xi,yi] = meshgrid(xmin+0.5*xsmpres:xsmpres:xmax-0.5*xsmpres,ymin+0.5*ysmpres:ysmpres:ymax-0.5*ysmpres);

        cdstack_interp_dry = zeros(nnrows,nncols,cdslices);
        parfor n = 1:cdslices
           slicelist = reshape(cdstack_dry(:,n),[],1);
           newslice = griddata(lonlist,latlist,slicelist,xi,yi,'linear');
           %F = TriScatteredInterp(lonlist,latlist,slicelist,'linear');
           %newslicelist = F(xi,yi);
           %newslice = reshape(newslicelist,nnrows,nncols);
           cdstack_interp_dry(:,:,n)= flipud(newslice); % flipud due to ypositive downpage for matlab figs, and ypositive uppage for utmy
        end

        cdstack_interp_wet = zeros(nnrows,nncols,cdslices);
        parfor n = 1:cdslices
           slicelist = reshape(cdstack_wet(:,n),[],1);
           newslice = griddata(lonlist,latlist,slicelist,xi,yi,'linear');
           %F = TriScatteredInterp(lonlist,latlist,slicelist,'linear');
           %newslicelist = F(xi,yi);
           %newslice = reshape(newslicelist,nnrows,nncols);
           cdstack_interp_wet(:,:,n)= flipud(newslice); % flipud due to ypositive downpage for matlab figs, and ypositive uppage for utmy
        end
        % keeping the coordinates in the same grid as the data
        xi = flipud(xi);
        yi = flipud(yi);

        % Pull out delays from cdstack layers that match dem heights
        wetcorrection = ones(nnrows,nncols);
        hydrcorrection = ones(nnrows,nncols);
        rounddem = round(dem/vertres);
        rounddem(dem < 0)=0;

        % cope with the case that NaN are present. This can be sea-level
        rounddem(isnan(dem))=0;

        for i=1:nnrows
            for j=1:nncols
                wetcorrection(i,j) = cdstack_interp_wet(i,j,rounddem(i,j)+1);
            end
        end

        for i=1:nnrows
            for j=1:nncols
                hydrcorrection(i,j) = cdstack_interp_dry(i,j,rounddem(i,j)+1);
            end
        end

        % note that this is a one way Zenith delay and not a slant delay.
        % Units are in cm
        % Output wet correction
        wetcorrection = wetcorrection*100;                 % delay in [cm]
        fid = fopen(outfile,'w');
        data_write =  [reshape(xi,[],1) reshape(yi,[],1) reshape(wetcorrection,[],1)]';
        tally = fwrite(fid,data_write,'double');
        fclose(fid);
        fid = fopen(outfile2,'w');
        fwrite(fid,wetcorrection','real*4');
        fclose(fid);

        % Output dry correction
        hydrcorrection = hydrcorrection*100;                 % delay in [cm]
        fid = fopen(hydroutfile,'w');
        data_write =  [reshape(xi,[],1) reshape(yi,[],1) reshape(hydrcorrection,[],1)]';
        tally = fwrite(fid,data_write,'double');
        fclose(fid);
        fid = fopen(hydroutfile2,'w');
        fwrite(fid,hydrcorrection','real*4');
        fclose(fid);

        fprintf([num2str(d) ' completed out of ' num2str(n_dates) '\n' ])

    else
        fprintf(['No wrf data for ' date ' \n'])
    end
end


%                           solvePhaseMax.m
%
%  Implementation of the PhaseMax algorithm proposed in the paper using
%  FASTA. Note: For this code to run, the solver "fasta.m" must be in your
%  path.
%
%% I/O
%  Inputs:
%     A:    m x n matrix or a function handle to a method that
%           returns A*x.
%     At:   The adjoint (transpose) of 'A'. If 'A' is a function handle, 'At'
%           must be provided.
%     b0:   m x 1 real,non-negative vector consists of all the measurements.
%     x0:   n x 1 vector. It is the initial guess of the unknown signal x.
%     opts: A struct consists of the options for the algorithm. For details,
%           see header in solvePhaseRetrieval.m or the User Guide.
%
%     Note: When a function handle is used, the
%     value of 'At' (a function handle for the adjoint of 'A') must be
%     supplied.
%
%  Outputs:
%     sol:  n x 1 vector. It is the estimated signal.
%     outs: A struct consists of the convergence info. For details,
%           see header in solvePhaseRetrieval.m or the User Guide.
%
%
%  See the script 'testPhaseMaxGaussian.m' and 'testPhaseMaxFourier.m' for
%  two examples of proper usage of this function.
%
%% Notations
%
%
%% Algorithm Description
%  Solve the PhaseMax signal reconstruction problem
%         maximize <x0,x>
%         subject to |Ax|<=b0
%
%  The problem is solved by approximately enforcing the constraints using a
%  quadratic barrier function.  A continuation method is used to increase
%  the strength of the barrier function until a high level of numerical
%  accuracy is reached.
%  The objective plus the quadratic barrier has the form
%     <-x0,x> + 0.5*max{|Ax|-b,0}^2.
%
%  For a detailed explanation, see the PhaseMax paper referenced below. For
%  more details about FASTA, see the FASTA user guide, or the paper "A
%  field guide to forward-backward splitting with a FASTA implementation."
%
%% References
%  Paper Title:   PhaseMax: Convex Phase Retrieval via Basis Pursuit
%  Authors:       Tom Goldstein, Christoph Studer
%  arXiv Address: https://arxiv.org/abs/1610.07531
%
%  Copyright Goldstein & Studer, 2016.  For more details, visit
%  https://www.cs.umd.edu/~tomg/projects/phasemax/

function [sol, outs] = solvePhaseMax(A,At,b0,x0,opts)
    % Initialization
    m = length(b0);              % number of measurements
    n = length(x0);              % length of the unknown signal
    remainIters = opts.maxIters; % The remaining fasta iterations we have.
                                 % It's initialized to opts.maxIters.

    %  Normalize the initial guess relative to the number of measurements
    x0 = (x0/norm(x0(:)))*mean(b0(:))*(m/n)*100;

    %  re-scale the initial guess so that it approximately satisfies |Ax|=b
    sol = x0.* min(b0./abs(A(x0)));

    % Initialize values potentially computed at each round.
    ending = 0;       % Indicate whether any of the ending condition (except maxIters) has been met in FASTA.
    iter = 0;
    currentTime = [];
    currentResid = [];
    currentReconError = [];
    currentMeasurementError = [];

    % Initialize vectors for recording convergence information
    [solveTimes,measurementErrors,reconErrors,residuals] = initializeContainers(opts);

    %%  Define objective function components for the gradient descent method FASTA
    f = @(z) 0.5*norm(max(abs(z)-b0,0))^2;     % f(z) = 0.5*max{|x|-b,0}^2 : This is the quadratic penalty function
    gradf = @(z)  (sign(z).*max(abs(z)-b0,0)); % The gradient of the quadratic penalty

    % Options to hand to fasta
    fastaOpts.maxIters = opts.maxIters;
    fastaOpts.stopNow = @(x, iter, resid, normResid, maxResid, opts) ...
        processIteration(x, resid);  % Use customized stopNow in order to get
                                         % solveTime, residual and error at each iteration.
    fastaOpts.verbose=0;
    startTime = tic;                    % Start timer
    constraintError = norm(abs(A(sol))-b0);  % Keep track of the current error in the solution.
    while remainIters > 0 & ~ending     % Iterate over continuation steps
        g = @(x) -real(x0'*x);          % The linear part of the objective
        proxg = @(x,t) x+t*x0;          % The proximal operator of the linear objective
        fastaOpts.tol = norm(x0)/100;   % use a tighter tolerance when the solution is more exact
        % Call FASTA to solve the inner minimization problem
        [sol, fastaOuts] = fasta(A, At, f, gradf, g, proxg, sol, fastaOpts);  % Call a solver to minimize the quadratic barrier problem

        fastaOpts.tau = fastaOuts.stepsizes(end);     % Record the most recent stepsize for recycling.
        x0 = x0/10;                                   % do continuation - this makes the quadratic penalty stronger

        % Update the max number of iterations for fasta
        remainIters = remainIters - fastaOuts.iterationCount;
        fastaOpts.maxIters = min(opts.maxIters, remainIters);

        % Monitor convergence and check stopping conditions
        newConstraintError = norm(max(abs(A(sol))-b0,0));
        relativeChange = abs(constraintError-newConstraintError)/norm(b0);
        if relativeChange <opts.tol   % terminate when error reduction stalls
            break;
        end
        constraintError = newConstraintError;
    end

    % Create output according to the options chosen by user
    outs = generateOutputs(opts, iter, solveTimes, measurementErrors, reconErrors, residuals);

    % Display verbose output if specified
    if opts.verbose == 1
        displayVerboseOutput(iter, currentTime, currentResid, currentReconError, currentMeasurementError);
    end

    % Runs code upon each FASTA iteration. Returns whether FASTA should
    % terminate.
    function stop = processIteration(x, residual)
        iter = iter + 1;
        % Record convergence information and check stopping condition
        % If xt is provided, reconstruction error will be computed and used for stopping
        % condition. Otherwise, residual will be computed and used for stopping
        % condition.
        if ~isempty(opts.xt)
            xt = opts.xt;
            %  Compute optimal rotation
            alpha = (x(:)'*xt(:))/(x(:)'*x(:));
            x = alpha*x;
            currentReconError = norm(x-xt)/norm(xt);
            if opts.recordReconErrors
                reconErrors(iter) = currentReconError;
            end
        end

        if isempty(opts.xt)
            currentResid = residual;
        end

        if opts.recordResiduals
            residuals(iter) = residual;
        end

        currentTime = toc(startTime);                % Record elapsed time so far
        if opts.recordTimes
            solveTimes(iter) = currentTime;
        end
        if opts.recordMeasurementErrors
            currentMeasurementError = norm(abs(A(sol)) - b0) / norm(b0);
            measurementErrors(iter) = currentMeasurementError;
        end

        % Display verbose output if specified
        if opts.verbose == 2
            displayVerboseOutput(iter, currentTime, currentResid, currentReconError, currentMeasurementError);
        end

        % Test stopping criteria.
        stop = false;
        if currentTime >= opts.maxTime % Stop if we're run over the max runtime
            stop = true;
        end
        if ~isempty(opts.xt) % If true solution is specified, terminate when close to true solution
            assert(~isempty(currentReconError),'If xt is provided, currentReconError must be provided.');
            stop = stop || currentReconError < opts.tol;
            ending = stop;  % When true, this flag will terminate outer loop
        end
        stop = stop || residual < fastaOpts.tol;    % Stop FASTA is the tolerance is reached
    end

end