% Copyright (c) 2010 Kristopher L. Kuhlman (klkuhlm at sandia dot gov)
% 
% Permission is hereby granted, free of charge, to any person obtaining a copy
% of this software and associated documentation files (the "Software"), to deal
% in the Software without restriction, including without limitation the rights
% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
% copies of the Software, and to permit persons to whom the Software is
% furnished to do so, subject to the following conditions:
% 
% The above copyright notice and this permission notice shall be included in
% all copies or substantial portions of the Software.
% 
% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
% OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
% THE SOFTWARE.
%
% this function is part of the set of MATLAB scripts 
% comprising the implementation of a universal kriging
% program that computes estimates of potential via
% cokriging, taking into account boundary conditions.
%
% Kuhlman, K.L., and E. Pardo-Igúzquiza, 2010. Universal
% cokriging of hydraulic heads accounting for boundary 
% conditions, Journal of Hydrology, 384(1–2), 14–25.
% http://dx.doi.org/10.1016/j.jhydrol.2010.01.002

clear;
nz = 40;

wclocs = [-3/4, -1/2];  % 1/2 unit south of line of observations
wdata =  [3/5,  4/5; 3/5, 4/5];    % arbitrary unit vectors

% ################################################
sgn = sum(sign(wdata),2);

% if vector is on x or y axis, or in quadrants 1 or 3,
% just take abs() of entire vector

% if vector is in Q2 or Q4, only take abs of x part
q2q4 = (sgn == 0);  % <- logical mask for Q2/Q4 membership

% move all negative x-axis and Q3 vectors up to Q1
wdata(~q2q4, 1:2) = abs(wdata(~q2q4, 1:2));

% move all Q2 vectors down to Q4
wdata(q2q4, 1) = abs(wdata(q2q4, 1));
% ################################################    

ntypes = 3;
ncovar = 3;
s = 7;

nrho = 15;
rhov = logspace(-2,0,nrho);
C = zeros(nz,ntypes,nrho,ncovar);
D = zeros(nz,ntypes,nrho,ncovar);

% parameters in models (not used in Gaussian)
m.alpha = 1.5;
m.nu = 2.25;

m.nug = 0.01;
m.sigsq = 13.0;
m.a = 6.0;

dum = [-999.0, 999.0];

zlocs = [linspace(-3,3,nz)', zeros(nz,1)];
hvec = zlocs(1:nz,1:2) - wclocs(ones(nz,1),1:2);
h = sqrt(dot(hvec,hvec,2));

for kk = 1:ncovar
    
    m.type = kk;
    for jj = 1:nrho
        
        m.rho = rhov(jj);
        % for fd approximation of gradient vector at single point
        wlocs = cat(3, wclocs + m.rho*wdata(1,:), wclocs - m.rho*wdata(1,:));
        
        % for fd approximation of gradient vector at row of zlocs
        wzlocs = cat(3,zlocs + m.rho*wdata(ones(nz,1)*2,:), zlocs - m.rho*wdata(ones(nz,1)*2,:));        
        
%         if kk == 1 && jj == s
%             % plot some stuff for visualizing test data
%             % #############################################################
%             figure(10)
%             subplot(211)
%             plot(zlocs(:,1),zlocs(:,2),'k*')
%             hold on
%             % h vectors
%             quiver(wclocs(ones(nz,1),1),wclocs(ones(nz,1),2),hvec(:,1),hvec(:,2), ...
%                 'g:','LineWidth',2,'Autoscale','off')
%             plot(wclocs(1,1),wclocs(1,2),'b*')
%             % boundary vectors
%             quiver(wclocs(1,1),wclocs(1,2),wdata(1,1),wdata(1,2), ...
%                 'b','LineWidth',2,'Autoscale','off')
%             % projections
%             hproj1 = dot(hvec(1:nz,1:2),wdata(ones(nz,1),1:2),2);
%             quiver(zlocs(:,1),zlocs(:,2),hproj1(:)*wdata(1,1),hproj1(:)*wdata(1,2), ...
%                 'r','LineWidth',2,'Autoscale','off')
%             daspect([1,1,1])
% %             w = axis();
%             title('first deriv true vector setup')
%             hold off
%             
%             subplot(212)
%             plot(zlocs(:,1),zlocs(:,2),'k*')
%             hold on
%             plot(wclocs(1,1),wclocs(1,2),'b*')
%             plot(wlocs(:,1,1),wlocs(:,2,1),'g.')
%             plot(wlocs(:,1,2),wlocs(:,2,2),'g.')
%             colors = {'b-','c-'};
%             for v = 1:nz
%                 for src = 1:2
%                     plot([wlocs(1,1,src),zlocs(v,1)],[wlocs(1,2,src),zlocs(v,2)],colors{src})
%                 end
%             end
%             daspect([1,1,1])
% %             axis(w)
%             title('first deriv approximate vector setup')
%             hold off
%             
%             figure(11)
%             subplot(211)
%             plot(zlocs(:,1),zlocs(:,2),'k*')
%             hold on
%             hunit = zeros(nz,2);
%             for v = 1:nz
%                 hunit(v,:) = hvec(v,:)./norm(hvec(v,:));
%             end
%             quiver(wclocs(ones(nz,1),1),wclocs(ones(nz,1),2),hvec(:,1),hvec(:,2), ...
%                 'g:','LineWidth',1,'Autoscale','off')
%             quiver(wclocs(ones(nz,1),1),wclocs(ones(nz,1),2),hunit(:,1),hunit(:,2), ...
%                 'g','LineWidth',2,'Autoscale','off')
%             plot(wclocs(:,1),wclocs(:,2),'b*')
%             quiver(wclocs(1,1),wclocs(1,2),wdata(1,1),wdata(1,2),'b','LineWidth',2)
%             quiver(zlocs(:,1),zlocs(:,2),wdata(ones(nz,1)*2,1),wdata(ones(nz,1)*2,2), ...
%                 'c','LineWidth',2,'Autoscale','off')
%             quiver(zlocs(:,1),zlocs(:,2),hproj1(:)*wdata(1,1),hproj1(:)*wdata(1,2), ...
%                 'r','LineWidth',1,'Autoscale','off')
%             hproj2 = dot(hvec(1:nz,1:2),wdata(ones(nz,1)*2,1:2),2);
%             quiver(zlocs(:,1),zlocs(:,2),hproj2(:)*wdata(2,1),hproj2(:)*wdata(2,2), ...
%                 'm','LineWidth',2,'Autoscale','off')
%             daspect([1,1,1])
% %             w = axis();
%             title('second deriv true vector setup')
%             hold off
%             
%             subplot(212)
%             plot(zlocs(:,1),zlocs(:,2),'k*')
%             hold on
%             plot(wclocs(:,1),wclocs(:,2),'bo')
%             plot(wlocs(:,1,1),wlocs(:,2,1),'g.')
%             plot(wlocs(:,1,2),wlocs(:,2,2),'g.')
%             plot(wzlocs(:,1,1),wzlocs(:,2,1),'r.')
%             plot(wzlocs(:,1,2),wzlocs(:,2,2),'r.')
%             colors = {'b-','c-','r-','m-'};
%             for v = 1:nz
%                 for src = 1:2
%                     for trg = 1:2
%                         if src == 1
%                             idx = trg;
%                         else
%                             idx = trg + 2;
%                         end
%                         plot([wlocs(1,1,src),wzlocs(v,1,trg)],[wlocs(1,2,src),wzlocs(v,2,trg)],colors{idx})
%                     end
%                 end
%             end
%             daspect([1,1,1])
% %             axis(w)
%             title('second deriv approximate vector setup')
%             hold off
%         end
        
        for j=1:nz
            % function not for checking but for plotting
            C(j,1,jj,m.type) = covar_derivs(m,0, zlocs(j,:), wclocs(1,:), dum, dum);
            D(j,1,jj,m.type) = C(j,1,jj,m.type);
            
            % h is a vector point from z2 to z1 (here from wclocs to each of nz zlocs)
            % e1 is a unit vector corresponding to the boundary
            % condition to be enforced (at the point z1).
            C(j,2,jj,m.type) = covar_derivs(m,1, zlocs(j,:), wclocs(1,:), wdata(1,:), dum);
            D(j,2,jj,m.type) = (covar_fcns(m, zlocs(j,1:2), wlocs(1,1:2,1),m.nug) - ...
                                covar_fcns(m, zlocs(j,1:2), wlocs(1,1:2,2),m.nug))./(2*m.rho);
            
            % h points from z2 to z1; e1 is the vector at z1 and e2 is the
            % vector at point z2.  wdata(1) is associated with point z2,
            % wdata(2) is associated with point z1.
            C(j,3,jj,m.type) = covar_derivs(m, 2, zlocs(j,:), wclocs(1,:), wdata(1,:), wdata(2,:));
            D(j,3,jj,m.type) = (covar_fcns(m, wzlocs(j,1:2,1), wlocs(1,1:2,1),m.nug) - ...
                                covar_fcns(m, wzlocs(j,1:2,2), wlocs(1,1:2,1),m.nug) - ...
                                covar_fcns(m, wzlocs(j,1:2,1), wlocs(1,1:2,2),m.nug) + ...
                                covar_fcns(m, wzlocs(j,1:2,2), wlocs(1,1:2,2),m.nug))./(2*m.rho)^2;                            
        end
    end
        
    titles ={'C(h)','d/dx C(h)','d^2/dx^2 C(h)'};
    models = {'Gaussian','Generalized Cauchy','Matern'};
        
    if kk==3
        figure(kk)
        for i=1:ntypes
            subplot(2,2,i)
            plot(h,C(:,i,s,kk),'-g', h,D(:,i,s,kk),'*r')
            title([titles{i},' for ',models{kk}])
        end
    end

end

colors = {'r','g','b'};
figure()
for kk=1:3
    loglog(rhov,squeeze(sum(abs(C(:,2,:,kk) - D(:,2,:,kk)),1))/nz,['-',colors{kk}])
    hold on
    loglog(rhov,squeeze(sum(abs(C(:,3,:,kk) - D(:,3,:,kk)),1))/nz,[':',colors{kk}])
    xlabel('\rho')
    ylabel('absolute difference between true & FD')
    title('red=Gaussian, green=gen Cauchy, blue=Matern')
    hold on
end

figure()
for kk=1:3
loglog(rhov,abs(squeeze(sum(C(:,2,:,kk) - D(:,2,:,kk),1))/nz./squeeze(C(1,2,:,kk))),['-',colors{kk}])
hold on
loglog(rhov,abs(squeeze(sum(C(:,3,:,kk) - D(:,3,:,kk),1))/nz./squeeze(C(1,2,:,kk))),[':',colors{kk}])
xlabel('\rho')
ylabel('relative difference (true - FD)/true')
title('red=Gaussian, green=gen Cauchy, blue=Matern')
hold on
end