% Steven Callahan
% Math 6790
% Final Project

clear all;

truncated_p = input('Generate analysis for entire domain? (y/n)>','s');
if truncated_p == 'y'
  
  % read entire file
  load xbfile_full.dat;
  
  % read data into individual arrays
  xb_lat_full = xbfile_full(:,1);
  xb_lon_full = xbfile_full(:,2);
  xb_elev_full = xbfile_full(:,3);
  xb_full = xbfile_full(:,4);
  % set plotting range
  xi_full = -117.0:0.05:-109.0;
  yi_full = 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_full,YI_full] = meshgrid(xi_full,yi_full);
  
  % plot background temperature over entire domain
  ZZI_full = griddata(xb_lon_full,xb_lat_full,xb_full,XI_full,YI_full);
  contourf(XI_full,YI_full,ZZI_full,v);
  title('Background Temperature Over Entire Domain'); 
  colorbar;
  pause;

  % plot terrain over entire domain
  ZI_full = griddata(xb_lon_full,xb_lat_full,xb_elev_full,XI_full,YI_full);
  contour(XI_full,YI_full,ZI_full,w);
  title('Background Temperature Over Entire Domain'); 
  pause;
  
  % 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>');
  for k=1:443
    po(k,k)=sigo;
  end
  
  % define background error covariance matrix using lat, lon, and elev
  sigb = input('Enter the background error variance> ');
  
  radius = 6370;
  rad = input('Enter the background error decorrelation scale>');
  rad2 = rad * rad;
  pb=sparse(7999,7999);
  for j=1:7999
    dx = radius .* cos(pi .*xb_lat_full ./180.) .* pi .* (xb_lon_full ...
						  - xb_lon_full(j)) ./ 180.;
    dy = radius .* pi .* (xb_lat_full - xb_lat_full(j)) ./ 180.;
    dz = xb_elev_full - xb_elev_full(j); 
    r2 = dx .* dx + dy .* dy + 0.02 .*(dz .* dz);
    pb(:,j) = sigb .* exp(-r2./rad2);
  end
  
  % define forward transform matrix
  h = sparse(443,7999);
  d = zeros(443,1);
  for  k=1:443
    dx = radius .* cos(pi .* yo_lat(k) ./ 180.) .* pi ...
	 .*(yo_lon(k)-xb_lon_full) ./ 180.;
    dy = radius .* pi .* (yo_lat(k) - xb_lat_full) ./ 180.;
    dist = sqrt(dx.*dx + dy.*dy);
    [min_val,min_index] = min(dist);
    h(k,min_index) = 1;
    d(k) = yo(k)-xb_full(min_index);
  end
  
  % 3Dvar conjugate gradient solution
  pbhpo = pb' * h' * inv(po) * d;
  pbt = pb' + pb' * h' * inv(po) * h * pb;
  tic
  x3dv = xb_full + pb * cgs(pbt,pbhpo,.001,200);
  cgs_time = toc
  
  % plot 3DVar with cgs
  ZZV = griddata(xb_lon_full,xb_lat_full,x3dv,XI_full,YI_full);
  contourf(XI,YI,ZZV,v);
  title('3DVar Analysis Over Entire Domain');
  colorbar;

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

  % plot background temperature over sub domain
  ZZI = griddata(xb_lon,xb_lat,xb,XI,YI);
  contourf(XI,YI,ZZI,v);
  title('Background Temperature Over Sub Domain');
  colorbar;
  pause;

  % plot terrain over sub domain
  ZI = griddata(xb_lon,xb_lat,xb_elev,XI,YI);
  contour(XI,YI,ZI,w);
  title('Terrain Over Sub Domain');
  pause;
  
  % 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>');
  for k=1:443
    po(k,k)=sigo;
  end
  
  % define background error covariance matrix using lat, lon, and elev
  sigb = input('Enter the background error variance> ');
  
  radius = 6370;
  rad = input('Enter the background error decorrelation scale>');
  rad2 = rad * rad;
  pb=zeros(4284,4284);
  for j=1:4284
    dx = radius .* cos(pi .*xb_lat ./180.) .* pi .* (xb_lon - xb_lon(j)) ...
	 ./ 180.;
    dy = radius .* pi .* (xb_lat - xb_lat(j)) ./ 180.;
    dz = xb_elev - xb_elev(j); 
    r2 = dx .* dx + dy .* dy + 0.02 .*(dz .* dz);
    pb(:,j) = sigb .* exp(-r2./rad2);
  end
  
  % define forward transform matrix
  h = zeros(443,4284);
  d = zeros(443,1);
  for  k=1:443
    dx = radius .* cos(pi .* yo_lat(k) ./ 180.) .* pi .*(yo_lon(k)-xb_lon) ...
	 ./ 180.;
    dy = radius .* pi .* (yo_lat(k) - xb_lat) ./ 180.;
    
    dist = sqrt(dx.*dx + dy.*dy);
    [min_val,min_index] = min(dist);
    h(k,min_index) = 1;
    d(k) = yo(k)-xb(min_index);
  end
  
  % 3Dvar conjugate gradient solution
  pbhpo = pb' * h' * inv(po) * d;
  pbt = pb' + pb' * h' * inv(po) * h * pb;
  tic
  x3dv = xb + pb * cgs(pbt,pbhpo,.001,200);
  cgs_time = toc
  
  % plot 3DVar with cgs
  ZZV = griddata(xb_lon,xb_lat,x3dv,XI,YI);
  contourf(XI,YI,ZZV,v);
  title('3DVar Analysis Over Sub Domain');
  colorbar;
  pause

  contourf(XI,YI,ZZV,v);
  hold on
  plot(yo_lon,yo_lat,'kx');
  title('3DVar Analysis Over Sub Domain with Observation Points');
  colorbar;
  hold off
  pause
  
  % 3Dvar generalized minimum residual solution
  tic
  x3dv = xb + pb * gmres(pbt,pbhpo,1000,0.001);
  gmres_time = toc
  
  % plot 3DVar with gmres
  ZZV = griddata(xb_lon,xb_lat,x3dv,XI,YI);
  contourf(XI,YI,ZZV,v);
  title('3DVar Analysis Over Sub Domain using GMRES');
  colorbar;
  hold off
  pause
  
  contourf(XI,YI,ZZV,v);
  hold on
  plot(yo_lon,yo_lat,'kx');
  title('3DVar Analysis Over Sub Domain using GMRES with Observation Points');
  colorbar;
  hold off
  pause
end
