close all;
clear all;

%read truncated file
load xbfile.dat;
% read data into individual arrays
xb_lat = xbfile(:,1);
xb_lon = xbfile(:,2);
xb_elev = xbfile(:,3);
xb = xbfile(:,4);
% set plotting range
xi = -117.0:0.05:-109.0;
yi = 39.0:0.05:44.0;
v = [-8 -6 -4 -2 0 2 4 6 8 10 12 14 16 18];
w = [250 500 750 1000 1250 1500 1750 2000 2250 2500 2750 3000 3250 3500 3750];
[XI,YI] = meshgrid(xi,yi);
% load observation data
[yo_lat,yo_lon,yo_elev,yo_elev2,yo,yo_nam,yo_nam2] = textread('merged.dat','%f %f %f %f %f %s %s');


% define observation error covariance matrix
sigo = input('Enter the observation error variance>');
No = length(yo);
po = zeros(No,No);
for k=1:No
    po(k,k)=sigo;
end
%transform to sparse matrix to decrease memory requirement and increase computation
%efficiency
po = sparse(po);


% define background error covariance matrix
sigb = input('Enter the background error variance>');
radius = 6370;
rad = input('Enter the background error decorrelation scale>');
%%% move the redundant calculations out of the for loop to increase
%%% computation efficiency
rad2 = rad^2;
radz2 = (rad/50)^2;
mrad2 =5*rad2;
mradz2 =5*radz2;
rca_xb_lat = radius * cos(pi * xb_lat/180.)*pi/180.0;
ra = radius * pi /180.0;
Nb = length(xb);
pb=zeros(Nb, Nb);
for j=1:Nb
    for i=1:Nb
        dx = rca_xb_lat(i) * (xb_lon(i)-xb_lon(j)); 
        dy = ra * (xb_lat(i) - xb_lat(j));
        dz2 = ((xb_elev(i)-xb_elev(j))/1000)^2;
        r2 = dx*dx + dy*dy;
        if(r2<mrad2 && dz2<mradz2)
            pb(i,j) =  sigb * exp (-r2/rad2 -dz2/radz2);
        end
    end
end
pb = sparse(pb);


% define forward transform matrix
h = zeros(No,Nb);
rca_yo_lat = radius * cos(pi * yo_lat/180.)*pi/180.0;
ra = radius * pi /180.0;
for  k=1:No
    min = 6370;
    loc = 1;    
    for i=1:Nb
        dx = rca_yo_lat(k) * (yo_lon(k)-xb_lon(i));
        dy = ra * (yo_lat(k) - xb_lat(i));
        dist = sqrt(dx*dx + dy*dy);
        if(dist<min)
            min = dist;
            loc = i;
        end
    end
    h(k,loc) = 1;
end
h=sparse(h);


% define the observation innovation matrix
d = yo - h*xb;


% 3Dvar conjugate gradient solution
pbhpo = pb' * h' * inv(po) * d;
pbt = pb' + pb' * h' * inv(po) * h * pb;

x3dv = xb + pb * cgs(pbt,pbhpo,1e-6,1000); 

ZZV = griddata(xb_lon,xb_lat,x3dv,XI,YI);
contourf(XI,YI,ZZV,v);
colorbar;
