def test_vrtdiv():
    path = ('/archive/Spencer.Hill/am2/am2clim_reyoi/gfdl.ncrc2-default-prod/'
            'pp/atmos_level/ts/monthly/1yr/atmos_level.198301-198312.')

    # Vertically defined, sigma levels.
    u_arr = xray.open_dataset(path + 'ucomp.nc').ucomp
    v_arr = xray.open_dataset(path + 'vcomp.nc').vcomp
    vort, divg = compute_vrtdiv(u_arr, v_arr)
    assert vort.shape == u_arr.shape
    assert divg.shape == u_arr.shape
    np.testing.assert_array_equal(u_arr.lat, vort.lat)
    np.testing.assert_array_equal(u_arr.lon, vort.lon)
    np.testing.assert_array_equal(u_arr.time, vort.time)
    np.testing.assert_array_equal(u_arr.pfull, vort.pfull)

    # Not vertically defined.
    u0 = u_arr[:,0]
    v0 = v_arr[:,0]
    vort0, divg0 = compute_vrtdiv(u0, v0)
    assert vort0.shape == u0.shape
    assert divg0.shape == u0.shape

    # Dummy case: zeros everywhere
    u_arr_zeros = xray.DataArray(np.zeros_like(u_arr.values), dims=u_arr.dims,
                                 coords=u_arr.coords)
    v_arr_zeros = u_arr_zeros.copy()
    vort_zeros, divg_zeros = compute_vrtdiv(u_arr_zeros, v_arr_zeros)
    assert not vort_zeros.any()
    assert not divg_zeros.any()

    def test_coordinates_encoding(self):
        def equals_latlon(obj):
            return obj == 'lat lon' or obj == 'lon lat'

        original = Dataset({'temp': ('x', [0, 1]), 'precip': ('x', [0, -1])},
                           {'lat': ('x', [2, 3]), 'lon': ('x', [4, 5])})
        with self.roundtrip(original) as actual:
            self.assertDatasetIdentical(actual, original)
        with create_tmp_file() as tmp_file:
            original.to_netcdf(tmp_file)
            with open_dataset(tmp_file, decode_coords=False) as ds:
                self.assertTrue(equals_latlon(ds['temp'].attrs['coordinates']))
                self.assertTrue(equals_latlon(ds['precip'].attrs['coordinates']))
                self.assertNotIn('coordinates', ds.attrs)
                self.assertNotIn('coordinates', ds['lat'].attrs)
                self.assertNotIn('coordinates', ds['lon'].attrs)

        modified = original.drop(['temp', 'precip'])
        with self.roundtrip(modified) as actual:
            self.assertDatasetIdentical(actual, modified)
        with create_tmp_file() as tmp_file:
            modified.to_netcdf(tmp_file)
            with open_dataset(tmp_file, decode_coords=False) as ds:
                self.assertTrue(equals_latlon(ds.attrs['coordinates']))
                self.assertNotIn('coordinates', ds['lat'].attrs)
                self.assertNotIn('coordinates', ds['lon'].attrs)

 def test_write_groups(self):
     data1 = create_test_data()
     data2 = data1 * 2
     with create_tmp_file() as tmp_file:
         data1.to_netcdf(tmp_file, group='data/1')
         data2.to_netcdf(tmp_file, group='data/2', mode='a')
         with open_dataset(tmp_file, group='data/1') as actual1:
             self.assertDatasetIdentical(data1, actual1)
         with open_dataset(tmp_file, group='data/2') as actual2:
             self.assertDatasetIdentical(data2, actual2)

 def test_dask_roundtrip(self):
     with create_tmp_file() as tmp:
         data = create_test_data()
         data.to_netcdf(tmp)
         chunks = {'dim1': 4, 'dim2': 4, 'dim3': 4, 'time': 10}
         with open_dataset(tmp, chunks=chunks) as dask_ds:
             self.assertDatasetIdentical(data, dask_ds)
             with create_tmp_file() as tmp2:
                 dask_ds.to_netcdf(tmp2)
                 with open_dataset(tmp2) as on_disk:
                     self.assertDatasetIdentical(data, on_disk)

 def test_open_dataset(self):
     original = Dataset({'foo': ('x', np.random.randn(10))})
     with create_tmp_file() as tmp:
         original.to_netcdf(tmp)
         with open_dataset(tmp, chunks={'x': 5}) as actual:
             self.assertIsInstance(actual.foo.variable.data, da.Array)
             self.assertEqual(actual.foo.variable.data.chunks, ((5, 5),))
             self.assertDatasetAllClose(original, actual)
         with open_dataset(tmp) as actual:
             self.assertIsInstance(actual.foo.variable.data, np.ndarray)
             self.assertDatasetAllClose(original, actual)

    def setUp(self):

        # netcdfs we'll use for input and check output against
        nc_isnobal_input = 'test/data/isnobal_input.nc'
        nc_isnobal_output = 'test/data/isnobal_output.nc'

        # connect to the virtual watershed
        self.vwc = default_vw_client()

        # load NetCDF inputs and outputs from test data
        self.input_dataset = open_dataset(nc_isnobal_input)
        self.output_dataset = open_dataset(nc_isnobal_output)

        # insert NetCDF test input to virtual watershed
        input_mr_name = 'webapp-testing-input'

        modelruns = self.vwc.modelrun_search()
        unittest_uuids = [r['Model Run UUID'] for r in modelruns.records
                          if r['Model Run Name'] == 'webapp-testing-input']

        for u in unittest_uuids:
            s = self.vwc.delete_modelrun(u)
            print "pre-test cleanup success on %s: %s" % (u, str(s))

        self.model_run_uuid = \
            self.vwc.initialize_modelrun(
                model_run_name=input_mr_name,
                description='test in vwplatform',
                researcher_name='Matt Turner',
                keywords='test,isnobal,webapp')

        self.vwc.upload(self.model_run_uuid, nc_isnobal_input)

        self.start_datetime = '2010-10-01 00:00:00'
        self.end_datetime = '2010-10-01 16:00:00'

        md = metadata_from_file(nc_isnobal_input, self.model_run_uuid,
                                self.model_run_uuid,
                                'test input for isnobal run',
                                'Dry Creek', 'Idaho', model_name='isnobal',
                                start_datetime=self.start_datetime,
                                end_datetime=self.end_datetime,
                                model_set='inputs', taxonomy='geoimage',
                                model_set_taxonomy='grid')
        # import ipdb; ipdb.set_trace()

        self.input_uuid = self.vwc.insert_metadata(md).text

        time.sleep(1)

    def test_engine(self):
        data = create_test_data()
        with self.assertRaisesRegexp(ValueError, 'unrecognized engine'):
            data.to_netcdf('foo.nc', engine='foobar')
        with self.assertRaisesRegexp(ValueError, 'invalid engine'):
            data.to_netcdf(engine='netcdf4')

        with create_tmp_file() as tmp_file:
            data.to_netcdf(tmp_file)
            with self.assertRaisesRegexp(ValueError, 'unrecognized engine'):
                open_dataset(tmp_file, engine='foobar')

        netcdf_bytes = data.to_netcdf()
        with self.assertRaisesRegexp(ValueError, 'can only read'):
            open_dataset(BytesIO(netcdf_bytes), engine='foobar')

	def _calc_anomalies( self, *args, **kwargs ):
		'''
		calculate absolute or relative anomalies given a NetCDF file
		of the Climatic Research Unit (CRU) Historical Time Series.
		'''
		import xray

		# handle modeled vs. historical
		if self.ar5_modeled != None and self.ar5_historical != None:
			# parse the input name for some file metadata HARDWIRED!
			output_naming_dict = DownscaleAR5.standardized_fn_to_vars( self.ar5_modeled )
			variable = output_naming_dict[ 'variable' ]

			# read in both modeled and historical
			ds = xray.open_dataset( self.ar5_modeled )
			ds = ds[ variable ]
			clim_ds = xray.open_dataset( self.ar5_historical )
			# climatology
			clim_ds = clim_ds.loc[ {'time':slice(self.climatology_begin,self.climatology_end)} ]
			climatology = clim_ds[ variable ].groupby( 'time.month' ).mean( 'time' )
			del clim_ds

		elif self.ar5_historical is not None and self.ar5_modeled is None:
			output_naming_dict = standardized_fn_to_vars( self.ar5_historical )
			variable = output_naming_dict[ 'variable' ]

			# read in historical
			ds = xray.open_dataset( self.ar5_historical )
			# climatology
			climatology = ds.loc[ {'time':slice(self.climatology_begin,self.climatology_end)} ]
			climatology = climatology[ variable ].groupby( 'time.month' ).mean( 'time' )

		else:
			NameError( 'ERROR: must have both ar5_modeled and ar5_historical, or just ar5_historical' )

		if self.plev is not None:
			plevel, = np.where( ds.plev == self.plev )
			ds = ds[ :, plevel[0], ... ]
			climatology = climatology[ :, plevel[0], ... ]

		# anomalies
		if self.absolute == True:
			anomalies = ds.groupby( 'time.month' ) - climatology
		elif self.absolute == False:
			anomalies = ds.groupby( 'time.month' ) / climatology
		else:
			AttributeError( '_calc_anomalies (ar5): absolute can only be True or False' )
		return anomalies

 def roundtrip(self, data, **kwargs):
     f, tmp_file = tempfile.mkstemp(suffix='.nc')
     os.close(f)
     data.dump(tmp_file)
     roundtrip_data = open_dataset(tmp_file, **kwargs)
     os.remove(tmp_file)
     return roundtrip_data

def concat_to_nc( filelist, output_filename, dim='time', begin_time=None, end_time=None, nc_format='NETCDF4', **kwargs ):
	'''
	take list of consecutive netcdf files (made for CMIP5 data) and stack them into a 
	single larger netcdf file.  This was necessary to overcome some bugginess in how 
	MFDataset is dealing with different calendar units on different files.  This is 
	technically valid CF-Compliant metadata, but is tricky to work with.  This hack allows
	us to get around some of this unpredictable behavior.

	PARAMETERS:
	-----------
	filelist = [list] list of string file paths to the sorted netcdf files to stack together
	output_filename = [str] path to and name of the output file to be generated (.nc extension)
	dim = [str] dimension to stack on -- default is 'time'
	begin_time = [str] PANDAS style datetime string syntax -- used in xray
	end_time = [str] PANDAS style datetime string syntax -- used in xray
	format = [str] output NetCDF format desired. valid strings are:
					'NETCDF4', 'NETCDF4_CLASSIC', 'NETCDF3_64BIT', 'NETCDF3_CLASSIC'
					default is 'NETCDF4'
	**kwargs -- potential future arguments or overloaded args to pass through (none implemented)

	RETURNS:
	--------

	output_filename as string, with the important side-effect of writing data to disk

	'''
	import xray
	with xray.concat([ xray.open_dataset( i ).load() for i in filelist ], dim ) as ds:
		# time slicer condition
		if begin_time != None and end_time != None:
			ds = ds.loc[ { dim:slice( begin_time, end_time ) } ]
		if os.path.exists( output_filename ):
			os.remove( output_filename )
		ds.to_netcdf( output_filename, mode='w', format=nc_format )
	return output_filename

def get_sea_mask(ds):
    raw_mask = xray.open_dataset('mask.nc')
    # extract places where the nearest latitude or longitude (before or after)
    # is in the ocean
    sea_mask = ((raw_mask.reindex_like(ds, method='pad').sftlf < 100)
                & (raw_mask.reindex_like(ds, method='backfill').sftlf < 100))
    return sea_mask

def save_gpcp_years(datafile, savedir, yearmin=1997, yearmax=2015):
    """Read GPCP daily data and save to individual file for each year."""

    savestr = savedir + '/gpcp_daily_%d.nc'

    with xray.open_dataset(datafile) as ds:
        pcp = ds['PREC']
        dates = ds['yyyyddd']
        yy = dates // 1000
        ddd = dates - 1000 * yy
        pcp.coords['year'] = yy
        pcp.coords['day'] = ddd

        years = np.arange(yearmin, yearmax + 1)
        for year in years:
            print(year)
            ind = np.where(yy.values == year)[0]
            precip = pcp[ind]
            days = precip['day'].values
            precip = precip.drop(['year', 'day'])
            precip = precip.rename({'time' : 'day'})
            precip['day'] = days
            savefile = savestr % year
            print('Saving to ' + savefile)
            atm.save_nc(savefile, precip)

    return None

    def do_netcdf_load(self, buf):
        from cStringIO import StringIO
        import xray

        f = StringIO(buf)

        return xray.open_dataset(f)

    def test_mask_and_scale(self):
        f, tmp_file = tempfile.mkstemp(suffix='.nc')
        os.close(f)

        nc = nc4.Dataset(tmp_file, mode='w')
        nc.createDimension('t', 5)
        nc.createVariable('x', 'int16', ('t',), fill_value=-1)
        v = nc.variables['x']
        v.set_auto_maskandscale(False)
        v.add_offset = 10
        v.scale_factor = 0.1
        v[:] = np.array([-1, -1, 0, 1, 2])
        nc.close()

        # first make sure netCDF4 reads the masked and scaled data correctly
        nc = nc4.Dataset(tmp_file, mode='r')
        expected = np.ma.array([-1, -1, 10, 10.1, 10.2],
                               mask=[True, True, False, False, False])
        actual = nc.variables['x'][:]
        self.assertArrayEqual(expected, actual)

        # now check xray
        ds = open_dataset(tmp_file)
        expected = create_masked_and_scaled_data()
        self.assertDatasetEqual(expected, ds)
        os.remove(tmp_file)

    def test_dump_and_open_encodings(self):
        # Create a netCDF file with explicit time units
        # and make sure it makes it into the encodings
        # and survives a round trip
        f, tmp_file = tempfile.mkstemp(suffix='.nc')
        os.close(f)

        ds = nc4.Dataset(tmp_file, 'w')
        ds.createDimension('time', size=10)
        ds.createVariable('time', np.int32, dimensions=('time',))
        units = 'days since 1999-01-01'
        ds.variables['time'].setncattr('units', units)
        ds.variables['time'][:] = np.arange(10) + 4
        ds.close()

        xray_dataset = open_dataset(tmp_file)
        os.remove(tmp_file)
        xray_dataset.dump(tmp_file)

        ds = nc4.Dataset(tmp_file, 'r')

        self.assertEqual(ds.variables['time'].getncattr('units'), units)
        self.assertArrayEqual(ds.variables['time'], np.arange(10) + 4)

        ds.close()
        os.remove(tmp_file)

    def test_open_encodings(self):
        # Create a netCDF file with explicit time units
        # and make sure it makes it into the encodings
        # and survives a round trip
        f, tmp_file = tempfile.mkstemp(suffix='.nc')
        os.close(f)

        ds = nc4.Dataset(tmp_file, 'w')
        ds.createDimension('time', size=10)
        ds.createVariable('time', np.int32, dimensions=('time',))
        units = 'days since 1999-01-01'
        ds.variables['time'].setncattr('units', units)
        ds.variables['time'][:] = np.arange(10) + 4
        ds.close()

        expected = Dataset()

        time = pd.date_range('1999-01-05', periods=10)
        encoding = {'units': units, 'dtype': np.dtype('int32')}
        expected['time'] = ('time', time, {}, encoding)

        actual = open_dataset(tmp_file)

        self.assertXArrayEqual(actual['time'], expected['time'])
        actual_encoding = {k: v for k, v in actual['time'].encoding.iteritems()
                           if k in expected['time'].encoding}
        self.assertDictEqual(actual_encoding, expected['time'].encoding)

        os.remove(tmp_file)

    def __init__(self, imp):
        self.ROOTDIR = os.path.join(os.environ['HOME'],'Studies/Masters/Myroms/Msc_idealized/')
        self.roms = RunSetup( os.path.join( self.ROOTDIR, "experiments.yaml"), imp )

        grdfile = self.roms.retrive_grid('grd')
        
        self.Mr, self.Lr = grdfile['lat_rho'].shape # Number of J/I-direction INTERIOR RHO-points
        self.Mu, self.Lu = grdfile['lat_u'].shape   # Number of J/I-direction U-points
        self.Mv, self.Lv = grdfile['lat_v'].shape   # Number of J/I-direction V-points

        self.h  = grdfile['h']

        ################################################################################
        # Date/Time variables
        ################################################################################
        self.dstart = 0 
        tdays = (self.roms.ntimes*self.roms.dt)/(60*60*24)
        self.smstime = range(0, tdays+2)
        self.ndays = len(self.smstime)

        ################################################################################
        # Thermohaline field variables - SISPRES TS CLIMATOLOGY FIELDS
        ################################################################################
        lims = [-43.80, -42.00, -24.25, -22.65] 
        
        DATASETdir = '/Users/Phellipe/Studies/Masters/Datasets/'
        tsclim_file = '%sClimatology/rio_ts_climatology.nc'%DATASETdir
        self.tsclim = xray.open_dataset( tsclim_file )\
                        .sel(y=slice(lims[2],lims[3]), 
                                x=slice(lims[0], lims[1])) 
        
        self.x = self.tsclim.coords['x'].values
        self.y = self.tsclim.coords['y'].values

def load(var_name):
    path = ('/archive/Spencer.Hill/am2/am2clim_reyoi/gfdl.ncrc2-default-prod/'
            'pp/atmos_level/ts/monthly/1yr/atmos_level.198301-198312.')

    ds = xray.open_dataset(path + var_name + '.nc',
                           drop_variables=['nv', 'time_bounds'])
    return ds[var_name]

def load_dailyrel(datafiles, yearnm='year', onset_varnm='D_ONSET',
                  retreat_varnm='D_RETREAT'):

    ds = atm.load_concat(datafiles, concat_dim=yearnm)
    if isinstance(ds, xray.DataArray):
        ds = ds.to_dataset()
    varnms = ds.data_vars.keys()
    if onset_varnm is not None:
        onset = ds[onset_varnm]
        varnms.remove(onset_varnm)
    else:
        onset = np.nan * ds[yearnm]
    if retreat_varnm is not None:
        retreat = ds[retreat_varnm]
        varnms.remove(retreat_varnm)
    else:
        retreat = np.nan * ds[yearnm]

    # Remaining data variable is the data field
    varnm = varnms[0]
    data = ds[varnm]

    # Copy attributes from the first file in the list
    with xray.open_dataset(datafiles[0]) as ds0:
        data.attrs = ds0[varnm].attrs

    return data, onset, retreat

def load_netcdf(ncfile, group):
    """
    Load a ModVsObs object previously saved in a netcdf file
    """
    # Load the group into a ModVsObs object
    ds = xray.open_dataset(ncfile, group=group)

    # work out the varname
    varnames = ds.data_vars.keys()

    for vv in varnames:
        if 'mod' in vv:
            varname = vv.strip('_mod')

    # Load the two variables (as Pandas objects)
    TSobs = ds['%s_obs'%varname].to_pandas()
    TSmod = ds['%s_mod'%varname].to_pandas()

    # Load the attributes
    attrs = ds['%s_obs'%varname].attrs

    # Convert to a ModVsObs object
    # Put the data into a ModVsObs object (model first then observed)
    return ModVsObs(TSmod.index.to_pydatetime(),\
            TSmod.values,\
            TSobs.index.to_pydatetime(),\
            TSobs.values,\
            varname=varname,\
            long_name=attrs['long_name'], \
            units=attrs['units'], \
            stationid=group,\
        )