function grd = roms_get_grid(grd_file,scoord,tindex,calc_zuv)
% $Id: roms_get_grid.m 370 2008-07-16 12:56:22Z wilkin $
% grd = roms_get_grid(grd_file)
% grd = roms_get_grid(grd_file,scoord);
% grd = roms_get_grid(grd_file,outfile);
% grd = roms_get_grid(grd_file,outfile,zeta_input,calc_zuv);
%
% Gets the lon,lat,mask,depth [and z coordinates] from netcdf roms netcdf 
% grd_file or output file
% 
% Input:
%     grd_file: The roms netcdf grid file name
%           or, an existing grd structure to which the vertical coordinates 
%               are to be added or updated
%
% Optional inputs:
%     scoord:   ROMS his/rst/avg file from which the s-coord params can be
%               determined
%            or 4-element vector [theta_s theta_b Tcline N] where
%               Vtransform=1 and Vstretching=1 is assumed
%            or 6-element vector [theta_s theta_b Tcline N ...
%                                 Vtransform Vstretching]
%
%     zeta_in:  How to obtain zeta information to use
%               when including free surface height in calculating z-coordinates
%            0 implies assume zeta=0
%            integer implies use this time index into his/rst/avg
%               file to read zeta
%            2-d array of zeta values
%
%     calc_zuv: If present, this argument (any value) activates computing
%               the depths z_u and z_v on the u and v points of the 
%               ROMS C-grid
%            
% Output is a structure containing all the grid information
%
% John Wilkin
% Updated (Sept 2002) to correct scoordinate formulation and optionally
% include zeta in the calculation

if isstruct(grd_file)
  % if the first input is already a grd structure
  % the intention is to add vertical coordinates below
  grd = grd_file;
  
else
  % get the horizontal grid information from a ROMS grid file (or
  % history, average etc file)
  grd.grd_file = grd_file;
  
  varlist = ...
    { 'mask_rho','mask_psi','mask_u','mask_v',...
    'angle','h','pm','pn','f'};
  for v = varlist
    vname = char(v);
    try
      tmp = nc_varget(grd_file,vname);
      grd = setfield(grd,vname,tmp);
    catch
      warning(['Variable not found: ' vname])
    end
  end
  
  varlist = {'x_rho','y_rho','x_u','y_u','x_v','y_v','x_psi','y_psi'};
  if nc_isvar(grd_file,'x_rho')
    for v = varlist
      vname = char(v);
      try
        tmp = nc_varget(grd_file,vname);
        grd = setfield(grd,vname,tmp);
      catch
        warning(['Variable not found: ' vname])
      end
    end
  end
  
  varlist = ...
    { 'lon_rho','lat_rho','lon_psi','lat_psi',...
    'lon_v','lat_v','lon_u','lat_u'};
  for v = varlist
    vname = char(v);
    try
      tmp = nc_varget(grd_file,vname);
      grd = setfield(grd,vname,tmp);
    catch
      warning([vname ' not found. Substituting x/y coords instead'])
      if strcmp(vname(1:3),'lon')
        usevname = strrep(vname,'lon','x');
      else
        usevname = strrep(vname,'lat','y');
      end
      grd = setfield(grd,vname,getfield(grd,usevname));
    end
  end
    
  if isfield(grd,'mask_rho')
    grd.mask_rho_nan = grd.mask_rho;
    land = find(grd.mask_rho_nan==0);
    grd.mask_rho_nan(land) = NaN;
  else
    % there is no mask information in the file so create unit masks in case
    % code tries to use them
    grd.mask_rho = ones(size(grd.h));
    grd.mask_rho_nan = grd.mask_rho;
    grd.mask_u = ones(size(grd.h(:,2:end)));
    grd.mask_v = ones(size(grd.h(2:end,:)));
    grd.mask_psi = ones(size(grd.h(2:end,2:end)));
  end
  
  % If the grid file includes coastline data, such as a file being used
  % with the Rutgers version of editmask.m, load this too
  try
    grd.lon_coast = nc_varget(grd_file,'lon_coast');
  catch
  end
  try
    grd.lat_coast = nc_varget(grd_file,'lat_coast');
  catch
  end
  
end

if nargin > 1  
  
    h = grd.h;
    hmin=min(min(h));
    hmax=max(max(h));
    [Mp Lp]=size(h);
    L=Lp-1;
    M=Mp-1;
  
  % get z_r and z_w for the given s-coordinate parameters
  
  if ~ischar(scoord)
   
    theta_s = scoord(1);
    theta_b = scoord(2);
    Tcline  = scoord(3);
    N       = scoord(4);

    if (length(scoord) < 5),
      Vtransform = 1;
      Vstretching = 1;
    else
      Vtransform = scoord(5);
      Vstretching = scoord(6);
    end,
    hc = Tcline;
    
  else
  
    % input 'scoord' is the name of a his/avg/rst file name
    % so attempt to get s-coord params from the file
    
    if (nc_isvar(scoord,'Vtransform')),
      Vtransform = nc_varget(scoord,'Vtransform');
    else
      Vtransform = 1;
    end,
    if (nc_isvar(scoord,'Vstretching')),
      Vstretching = nc_varget(scoord,'Vstretching');
    else
      Vstretching = 1;
    end,
    
    theta_s = nc_varget(scoord,'theta_s');
    theta_b = nc_varget(scoord,'theta_b');
    Tcline  = nc_varget(scoord,'Tcline');
    N       = length(nc_varget(scoord,'Cs_r'));
    hc = nc_varget(scoord,'hc');
    if (Vtransform == 1),
      if isempty(hc),
        hc = min(h(:));
      end,
    end,
  
  end

  [s_rho,Cs_r]=stretching(Vstretching, theta_s, theta_b, hc, N, 0, 0);
  [s_w  ,Cs_w]=stretching(Vstretching, theta_s, theta_b, hc, N, 1, 0);

  Np = N+1;
  sc_r = s_rho;
  sc_w = s_w;

  % zeta  
  zeta = zeros(size(grd.h)); % default
  if nargin > 2 % option to include zeta in z calculation
    
    if tindex == 0 % if tindex==0 zeta defaults to zero
      % do nothing
      
    else % if tindex==0 zeta defaults to zero
      if length(tindex)==1
        % tindex is a single index to zeta in a roms output file
        if ~ischar(scoord)
          error([ 'Can''t process zeta from file in the case that ' ...
            ' scoord parameters are input as a vector'])'
        end
        zeta = nc_varget(scoord,'zeta',[tindex-1 0 0],[1 -1 -1]);
        if isempty(zeta)
          warning([ 'zeta not found in ' scoord '. Assuming zeta=0.'])
          zeta = zeros(size(grd.h));
        end
      else
        % tindex should be a 2-d field of zeta values
        if any(size(tindex)-size(grd.h))
          % sizes of zeta and h don't match
          error('input tindex as zeta does not match grid dimensions')
        else
          zeta = tindex;
        end
      end
    end
    
  end
  grd.zeta = zeta;
  
  % rho-points

  report = 0;
  
  h = grd.h;
  rgrid = 1;
  z_r = set_depth(Vtransform, Vstretching, theta_s, theta_b, hc, N, ...
                  rgrid, h, zeta, report);
  grd.z_r = reshape(z_r,[N size(h)]);
  clear z_r
  
  % w-points

  wgrid = 5;
  z_w = set_depth(Vtransform, Vstretching, theta_s, theta_b, hc, N, ...
                  wgrid, h, zeta, report);
  grd.z_w = reshape(z_w,[Np size(h)]);

  if nargin > 3
    
    % u-points (cell centres in vertical)

    ugrid = 3;
    z_u = set_depth(Vtransform, Vstretching, theta_s, theta_b, hc, N, ...
                    ugrid, h, zeta, report);
    grd.z_u = reshape(z_u,[N L Mp]);
    clear z_u
    
    % u-points (cell edges in vertical)

    z_uw = 0.5.*(z_w(1:L,1:Mp,:)+z_w(2:Lp,1:Mp,:));
    grd.z_uw = reshape(z_uw,[Np L Mp]);
    clear z_uw
    
    % v-points (cell centres in vertical)

    vgrid = 4;
    z_v = set_depth(Vtransform, Vstretching, theta_s, theta_b, hc, N, ...
                    vgrid, h, zeta, report);
    grd.z_v = reshape(z_v,[N Lp M]);
    clear z_v
    
    % v-points (cell edges in vertical)

    z_vw= 0.5.*(z_w(1:Lp,1:M,:)+z_w(1:Lp,2:Mp,:));
    
    grd.z_vw = reshape(z_vw,[Np Lp M]);
    clear z_w z_vw
    
  end 
  grd.theta_s = theta_s;
  grd.theta_b = theta_b;
  grd.Tcline = Tcline;
  grd.N = N;
  grd.hc = hc;
  grd.sc_w = sc_w;
  grd.Cs_w = Cs_w;
  grd.sc_r = sc_r;
  grd.Cs_r = Cs_r;
  grd.s_w = s_w;
  grd.s_rho = s_rho;
  
end
