# Path: atmos/data.py
# def load_concat(paths, var_ids=None, concat_dim='TIME', subset_dict=None,
#                 func=None, func_args=None, func_kw=None, squeeze=True, verbose=True):
#     """Load a variable from multiple files and concatenate into one.
# 
#     Especially useful for extracting variables split among multiple
#     OpenDAP files.
# 
#     Parameters
#     ----------
#     paths : list of strings
#         List of file paths or OpenDAP urls to process.
#     var_ids : str or list of str, optional
#         Name(s) of variable(s) to extract.  If None then all variables
#         are extracted and a Dataset is returned.
#     concat_dim : str
#         Name of dimension to concatenate along. If this dimension
#         doesn't exist in the input data, a new one is created.
#     subset_dict : dict of 2-tuples, optional
#         Dimensions and subsets to extract.  Each entry in subset_dict
#         is in the form {dim_name : (lower_or_list, upper)}, where:
#         - dim_name : string
#             Name of dimension to extract from.
#             The dimension name can be the actual dimension name
#             (e.g. 'XDim') or a generic name (e.g. 'lon') and get_coord()
#             is called to find the specific name.
#         - lower_or_list : scalar or list of int or float
#             If scalar, then used as the lower bound for the   subset range.
#             If list, then the subset matching the list will be extracted.
#         - upper : int, float, or None
#             Upper bound for subset range. If lower_or_list is a list,
#             then upper is ignored and should be set to None.
#     func : function, optional
#         Function to apply to each variable in each file before concatenating.
#         e.g. compute zonal mean. Takes one DataArray as first input parameter.
#     func_args : list, optional
#         List of numbered arguments to pass to func.
#     func_kw : dict or list of dict, optional
#         Dict of keyword arguments to pass to func. To use different values for
#         different files, make func_kw a list of the same length as the list of
#         file paths, with func_kw[i] containing a dict of keyword args for
#         path[i]. Otherwise, make func_kw a single dict to use for all paths.
#     squeeze : bool, optional
#         If True, squeeze out extra dimensions and add info to attributes.
#     verbose : bool, optional
#         If True, print updates while processing files.
# 
#     Returns:
#     --------
#     data : xray.DataArray or xray.Dataset
#         Data extracted from input files.
#     """
# 
#     # Number of times to attempt opening file (in case of server problems)
#     NMAX = 3
#     # Wait time (seconds) between attempts
#     WAIT = 5
# 
#     if var_ids is not None:
#         var_ids = utils.makelist(var_ids)
# 
#     def get_data(path, var_ids, subset_dict, func, func_args, func_kw):
#         with xray.open_dataset(path) as ds:
#             if var_ids is None:
#                 # All variables
#                 data = ds
#             else:
#                 # Extract specific variables
#                 data = ds[var_ids]
#             if subset_dict is not None:
#                 data = subset(data, subset_dict, copy=False)
#             if func is not None:
#                 data_out = xray.Dataset()
#                 if func_args is None:
#                     func_args = []
#                 if func_kw is None:
#                     func_kw = {}
#                 for nm in data.data_vars:
#                     vars_out = func(data[nm], *func_args, **func_kw)
#                     if not isinstance(vars_out, xray.Dataset):
#                         vars_out = vars_out.to_dataset()
#                     for nm2 in vars_out.data_vars:
#                         data_out[nm2] = vars_out[nm2]
#                 data = data_out
#             data.load()
#         return data
# 
#     pieces = []
#     func_kw = utils.makelist(func_kw)
#     paths = utils.makelist(paths)
#     if len(func_kw) == 1:
#         func_kw *= len(paths)
#     for p, kw in zip(paths, func_kw):
#         print_if(None, verbose, printfunc=disptime)
#         print_if('Loading ' + p, verbose)
#         attempt = 0
#         while attempt < NMAX:
#             try:
#                 piece = get_data(p, var_ids, subset_dict, func, func_args, kw)
#                 print_if('Appending data', verbose)
#                 pieces.append(piece)
#                 attempt = NMAX
#             except RuntimeError as err:
#                 attempt += 1
#                 if attempt < NMAX:
#                     print('Error reading file.  Attempting again in %d s' %
#                           WAIT)
#                     time.sleep(WAIT)
#                 else:
#                     raise err
# 
#     print_if('Concatenating data', verbose)
#     data = xray.concat(pieces, dim=concat_dim)
#     print_if(None, verbose, printfunc=disptime)
# 
#     if squeeze:
#         data = xr.squeeze(data)
# 
#     if len(data.data_vars) == 1:
#         # Convert from Dataset to DataArray for output
#         data = data[data.data_vars.keys()[0]]
# 
#     return data


import numpy as np
import matplotlib.pyplot as plt
import xray
import atmos.plots as ap
import atmos.data as dat
import atmos.utils as utils
from atmos.data import load_concat
"""Testing load_concat"""



# ----------------------------------------------------------------------
# Read data
url = ('http://goldsmr3.sci.gsfc.nasa.gov/opendap/MERRA_MONTHLY/'
'MAIMCPASM.5.2.0/2014/MERRA300.prod.assim.instM_3d_asm_Cp.2014%02d.hdf')
paths = [url % 1, url % 2, url % 3 ]

var_ids = ['U', 'V', 'T']

lon1, lon2 = 0, 100
lat1, lat2 = -45, 45

ds1 = load_concat(paths, var_ids)
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
ds2 = load_concat(paths, var_ids, subset_dict=subset_dict)
ds3 = load_concat(paths, 'PS', subset_dict=subset_dict)
u1 = load_concat(paths, 'U')

# Path: atmos/data.py
# def load_concat(paths, var_ids=None, concat_dim='TIME', subset_dict=None,
#                 func=None, func_args=None, func_kw=None, squeeze=True, verbose=True):
#     """Load a variable from multiple files and concatenate into one.
# 
#     Especially useful for extracting variables split among multiple
#     OpenDAP files.
# 
#     Parameters
#     ----------
#     paths : list of strings
#         List of file paths or OpenDAP urls to process.
#     var_ids : str or list of str, optional
#         Name(s) of variable(s) to extract.  If None then all variables
#         are extracted and a Dataset is returned.
#     concat_dim : str
#         Name of dimension to concatenate along. If this dimension
#         doesn't exist in the input data, a new one is created.
#     subset_dict : dict of 2-tuples, optional
#         Dimensions and subsets to extract.  Each entry in subset_dict
#         is in the form {dim_name : (lower_or_list, upper)}, where:
#         - dim_name : string
#             Name of dimension to extract from.
#             The dimension name can be the actual dimension name
#             (e.g. 'XDim') or a generic name (e.g. 'lon') and get_coord()
#             is called to find the specific name.
#         - lower_or_list : scalar or list of int or float
#             If scalar, then used as the lower bound for the   subset range.
#             If list, then the subset matching the list will be extracted.
#         - upper : int, float, or None
#             Upper bound for subset range. If lower_or_list is a list,
#             then upper is ignored and should be set to None.
#     func : function, optional
#         Function to apply to each variable in each file before concatenating.
#         e.g. compute zonal mean. Takes one DataArray as first input parameter.
#     func_args : list, optional
#         List of numbered arguments to pass to func.
#     func_kw : dict or list of dict, optional
#         Dict of keyword arguments to pass to func. To use different values for
#         different files, make func_kw a list of the same length as the list of
#         file paths, with func_kw[i] containing a dict of keyword args for
#         path[i]. Otherwise, make func_kw a single dict to use for all paths.
#     squeeze : bool, optional
#         If True, squeeze out extra dimensions and add info to attributes.
#     verbose : bool, optional
#         If True, print updates while processing files.
# 
#     Returns:
#     --------
#     data : xray.DataArray or xray.Dataset
#         Data extracted from input files.
#     """
# 
#     # Number of times to attempt opening file (in case of server problems)
#     NMAX = 3
#     # Wait time (seconds) between attempts
#     WAIT = 5
# 
#     if var_ids is not None:
#         var_ids = utils.makelist(var_ids)
# 
#     def get_data(path, var_ids, subset_dict, func, func_args, func_kw):
#         with xray.open_dataset(path) as ds:
#             if var_ids is None:
#                 # All variables
#                 data = ds
#             else:
#                 # Extract specific variables
#                 data = ds[var_ids]
#             if subset_dict is not None:
#                 data = subset(data, subset_dict, copy=False)
#             if func is not None:
#                 data_out = xray.Dataset()
#                 if func_args is None:
#                     func_args = []
#                 if func_kw is None:
#                     func_kw = {}
#                 for nm in data.data_vars:
#                     vars_out = func(data[nm], *func_args, **func_kw)
#                     if not isinstance(vars_out, xray.Dataset):
#                         vars_out = vars_out.to_dataset()
#                     for nm2 in vars_out.data_vars:
#                         data_out[nm2] = vars_out[nm2]
#                 data = data_out
#             data.load()
#         return data
# 
#     pieces = []
#     func_kw = utils.makelist(func_kw)
#     paths = utils.makelist(paths)
#     if len(func_kw) == 1:
#         func_kw *= len(paths)
#     for p, kw in zip(paths, func_kw):
#         print_if(None, verbose, printfunc=disptime)
#         print_if('Loading ' + p, verbose)
#         attempt = 0
#         while attempt < NMAX:
#             try:
#                 piece = get_data(p, var_ids, subset_dict, func, func_args, kw)
#                 print_if('Appending data', verbose)
#                 pieces.append(piece)
#                 attempt = NMAX
#             except RuntimeError as err:
#                 attempt += 1
#                 if attempt < NMAX:
#                     print('Error reading file.  Attempting again in %d s' %
#                           WAIT)
#                     time.sleep(WAIT)
#                 else:
#                     raise err
# 
#     print_if('Concatenating data', verbose)
#     data = xray.concat(pieces, dim=concat_dim)
#     print_if(None, verbose, printfunc=disptime)
# 
#     if squeeze:
#         data = xr.squeeze(data)
# 
#     if len(data.data_vars) == 1:
#         # Convert from Dataset to DataArray for output
#         data = data[data.data_vars.keys()[0]]
# 
#     return data


# Path: testing/testing-data-load_concat.py
import numpy as np
import matplotlib.pyplot as plt
import xray
import atmos.plots as ap
import atmos.data as dat
import atmos.utils as utils
from atmos.data import load_concat

"""Testing load_concat"""



# ----------------------------------------------------------------------
# Read data
url = ('http://goldsmr3.sci.gsfc.nasa.gov/opendap/MERRA_MONTHLY/'
'MAIMCPASM.5.2.0/2014/MERRA300.prod.assim.instM_3d_asm_Cp.2014%02d.hdf')
paths = [url % 1, url % 2, url % 3 ]

var_ids = ['U', 'V', 'T']

lon1, lon2 = 0, 100
lat1, lat2 = -45, 45

ds1 = load_concat(paths, var_ids)
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
ds2 = load_concat(paths, var_ids, subset_dict=subset_dict)
ds3 = load_concat(paths, 'PS', subset_dict=subset_dict)
u1 = load_concat(paths, 'U')

u2 = load_concat(paths[0], 'U')