% 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;

%% parameters to manually set 
% <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

% are boundary conditions being used? (1=true, 0=false)
bc_flag = 1;

% order of trend 
% (0=constant, 2=linear, 5=quadratic)
L = 2;

% direct head covariance model parameters
m.type = 3;  % model 1=Gaussian, 2=generalized Cauchy, 3=Matern

if m.type == 2
    m.alpha = 1.5;  % for Cauchy model
elseif m.type == 3
    m.nu = 2.2;  % for Matern model
end

m.a     = 6.35*0.4;   % range
m.sigsq = 13.0;   % sill
m.nug   = 0.01;    % head direct-covariance nugget

% output mesh related things
% nx = 56;  ny = 77;
% xmin = 601750.0;  ymin = 3566550.0;
% dx = 400.0; dy = 400.0;

nx = 20; ny = 20;
xmin = 0.0; ymin = 0.01;
dx = 0.499; dy = 0.499;
% >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
 
% ((((((((((((((((((((((((((((((((((((((((((((((((((
%% read / prepare head observations

z = load('test02_head.dat');  % 'test01_head.dat'
N = size(z,1);
d.zlocs = z(1:N,1:2);
d.zdata = z(1:N,3);  % col 3 is residual, col4 is fwh
% ))))))))))))))))))))))))))))))))))))))))))))))))))

% [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[
%%  read / prepare BC "observations"

if bc_flag == 1
    nfclocs = load('test01_nslocs.dat');   % no-flow boundary data
%    chclocs = load('test02_ewlocs.dat'); % constant head boundary data
      chclocs = [];
    d.wlocs = [nfclocs; chclocs];
    M = size(d.wlocs,1);
    
    
    % unit vector pointing normal to no-flow boundary
    % or parallel to specified head boundary
    nfdata = load('test01_nfvecs.dat');   % no-flow vectors
%    chdata = load('test02_nfvecs.dat');  % constant head vectors
     chdata = [];
    d.wdata = [nfdata;chdata];

    % ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^    
    % sanitize vector input
    % because unit vector orientation, but not +/- sense, is important
    % we therefore take  the Q1 or Q4 version of each vector for consisitency
    
    % values of sgn
    %----------------
    %         |
    %   Q2: 0 |  Q1: 2
    % --------+---------
    %   Q3:-2 |  Q4: 0
    %         |
    % axes are +/- 1
    
    sgn = sum(sign(d.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
    d.wdata(~q2q4, 1:2) = abs(d.wdata(~q2q4, 1:2));
    
    % move all Q2 vectors down to Q4
    d.wdata(q2q4, 1) = abs(d.wdata(q2q4, 1));   
    d.wdata(q2q4, 2) = -abs(d.wdata(q2q4, 2));

elseif bc_flag == 0
    M=0;
else
    error('illegal value for bc_flag: no bc=0 or with bc=1')
end
% ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]


%% uniform grid of calculation points
[X,Y] = meshgrid(linspace(xmin,xmin+nx*dx,nx), ...
                 linspace(ymin,ymin+ny*dy,ny));             
D = numel(X);

% observation points in a long x,y vector
d.estlocs(1:D,1:2) = [reshape(X,D,1),reshape(Y,D,1)];

% setup and solve Kriging matrices
% ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[z0,MSE] = solve_true_gradient_system(N,M,L,D,d,m,bc_flag);

% make output easier to plot
z0 = reshape(z0,ny,nx);
MSE = reshape(MSE,ny,nx);

%% plot
figure;
subplot(1,2,1)
contourf(X,Y,z0,20);
title('estimated head');
hold on;
plot(d.zlocs(:,1),d.zlocs(:,2),'k*');
if bc_flag == 1
    quiver(d.wlocs(:,1),d.wlocs(:,2),d.wdata(:,1),d.wdata(:,2),0.2);
end
colorbar;
daspect([1 1 1]);
hold off;

subplot(1,2,2)
contourf(X,Y,MSE,20);
title('MSE');
hold on;
plot(d.zlocs(:,1),d.zlocs(:,2),'k*');
if bc_flag == 1
    quiver(d.wlocs(:,1),d.wlocs(:,2),d.wdata(:,1),d.wdata(:,2),0.2);
end
colorbar;
daspect([1 1 1]);
hold off;
