content: """
  Replace corrupted data files with daily data re-downloaded with wget
  """
  
  import sys
  sys.path.append('/home/jwalker/dynamics/python/atmos-tools')
  sys.path.append('/home/jwalker/dynamics/python/atmos-read')
  
  import os
  import shutil
  import xarray as xray
  import numpy as np
  import collections
  import time
  import matplotlib.pyplot as plt
  import pandas as pd
  import atmos as atm
  import precipdat
  import merra
  
  # ----------------------------------------------------------------------
  datadir = '/net/eady/data1/jwalker/datastore/merra2/wget/'
  savedir = '/net/eady/data1/jwalker/datastore/merra2/merged/'
  probdata = pd.read_csv('scripts/merra_urls/merge_data.csv', index_col=0)
  
  # For each corrupted data file:
  # - load the corrupted data file
  # - load the new downloaded file for the problem day
  # - calculate d/dp and other stuff
  # - merge the data for the affected day
  # - save into data file for the year
  
  def latlon_filestr(lat1, lat2, lon1, lon2):
      """Return nicely formatted string for lat-lon range."""
      latstr = atm.latlon_str(lat1, lat2, 'lat')
      lonstr = atm.latlon_str(lon1, lon2, 'lon')
      return lonstr + '_' + latstr
  
  def latlon_data(var, lat1, lat2, lon1, lon2, plev=None):
      """Extract lat-lon subset of data."""
      name = var.name
      varnm = name
      subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
      latlonstr = latlon_filestr(lat1, lat2, lon1, lon2)
      if plev is not None:
          name = name + '%d' % plev
          subset_dict['plev'] = (plev, plev)
      var = atm.subset(var, subset_dict, copy=False, squeeze=True)
      var.name = name
      var.attrs['filestr'] = '%s_%s' % (name, latlonstr)
      var.attrs['varnm'] = varnm
      return var
  
  def pgradient(var, lat1, lat2, lon1, lon2, plev):
      """Return d/dp of a lat-lon variable."""
      pwidth = 100
      p1, p2 = plev - pwidth, plev + pwidth
      var = atm.subset(var, {'lat' : (lat1, lat2), 'lon' : (lon1, lon2),
                             'plev' : (p1, p2)}, copy=False, squeeze=True)
      latlonstr = latlon_filestr(lat1, lat2, lon1, lon2)
      attrs = var.attrs
      pname = atm.get_coord(var, 'plev', 'name')
      pdim = atm.get_coord(var, 'plev', 'dim')
      pres = var[pname]
      pres = atm.pres_convert(pres, pres.attrs['units'], 'Pa')
      dvar_dp = atm.gradient(var, pres, axis=pdim)
      dvar_dp = atm.subset(dvar_dp, {pname : (plev, plev)}, copy=False,
                           squeeze=True)
      varnm = 'D%sDP' % var.name
      name = '%s%d' % (varnm, plev)
      dvar_dp.name = name
      attrs['long_name'] = 'd/dp of ' + var.attrs['long_name']
      attrs['standard_name'] = 'd/dp of ' + var.attrs['standard_name']
      attrs['units'] = ('(%s)/Pa' % attrs['units'])
      attrs[pname] = plev
      attrs['filestr'] = '%s_%s' % (name, latlonstr)
      attrs['varnm'] = varnm
      dvar_dp.attrs = attrs
      return dvar_dp
  
  
  def var_calcs(filenm, varnm, plev, latlon=(-90, 90, 40, 120)):
      """Process a single variable from a single day."""
      lat1, lat2, lon1, lon2 = latlon
      if varnm == 'DUDP':
          nm, dp = 'U', True
      elif varnm == 'DOMEGADP':
          nm, dp = 'OMEGA', True
      else:
          nm, dp = varnm, False
      with xray.open_dataset(filenm) as ds:
          var = ds[nm].load()
      if dp:
          print('Computing d/dp')
          var = pgradient(var, lat1, lat2, lon1, lon2, plev)
      else:
          var = latlon_data(var, lat1, lat2, lon1, lon2, plev)
  
      return var
  
  def process_row(row, datadir, savedir):
      filenm1 = row['filename']
      year = row['year']
      varnm = row['varnm']
      plev = row['plev']
      jday = row['jday']
      filenm2 = datadir + row['datfile']
      savefile1 = filenm1
      savefile2 = savedir + os.path.split(filenm1)[1]
  
      print('%d, %s, plev=%d' % (year, varnm, plev))
      print('Reading original data from ' + filenm1)
      with xray.open_dataset(filenm1) as ds:
          var1 = ds[varnm].load()
      print('Processing new data from ' + filenm2)
      var2 = var_calcs(filenm2, varnm, plev)
  
      print('Merging data for jday %d' % jday)
      var = var1.copy()
      ind = jday - 1
      days = atm.get_coord(var1, 'day')
      if not days[ind] == jday:
          raise ValueError('Days not indexed from 1, need to edit code to handle')
      var[ind] = var2
      print('Saving to ' + savefile1)
      var.to_netcdf(savefile1)
      print('Saving to ' + savefile2)
      var.to_netcdf(savefile2)
  
      data = {'orig' : var1, 'new' : var2, 'merged' : var}
      return data
  
  # Make a copy of each of the original files -- only run this code once!
  # for filenm in probdata['filename']:
  #     shutil.copyfile(filenm, filenm.replace('.nc', '_orig.nc'))
  
  for i, row in probdata.iterrows():
      data = process_row(row, datadir, savedir)
  
  
  # Plot data to check
  def plot_data(probdata, savedir, i):
      row = probdata.iloc[i]
      filenm = row['filename']
      filenm = savedir + os.path.split(filenm)[1]
      jday = row['jday']
      varnm = row['varnm']
      with xray.open_dataset(filenm) as ds:
          var = ds[varnm].load()
      plt.figure(figsize=(16, 8))
      plt.suptitle(os.path.split(filenm)[1])
      plt.subplot(1, 3, 1)
      atm.pcolor_latlon(var.sel(day=(jday-1)))
      plt.title(jday - 1)
      plt.subplot(1, 3, 2)
      atm.pcolor_latlon(var.sel(day=jday))
      plt.title(jday)
      plt.subplot(1, 3, 3)
      atm.pcolor_latlon(var.sel(day=(jday+1)))
      plt.title(jday + 1)