# Path: atmos/data.py
# def get_ps_clim(lat, lon, datafile='data/topo/ncep2_ps.nc'):
#     """Return surface pressure climatology on selected lat-lon grid.
# 
#     Parameters
#     ----------
#     lat, lon : 1-D float array
#         Latitude and longitude grid to interpolate surface pressure
#         climatology onto.
#     datafile : string, optional
#         Name of file to read for surface pressure climatology.
# 
#     Returns
#     -------
#     ps : xray.DataArray
#         DataArray of surface pressure climatology interpolated onto
#         lat-lon grid.
#     """
# 
#     ds = ncload(datafile)
#     ps = ds['ps']
#     ps.attrs = utils.odict_insert(ps.attrs, 'title', ds.attrs['title'], pos=0)
# 
#     # Check what longitude convention is used in the surface pressure
#     # climatology and switch if necessary
#     lonmax = lon_convention(lon)
#     lon_ps = get_coord(ps, 'lon')
#     if lon_convention(lon_ps) != lonmax:
#         ps = set_lon(ps, lonmax)
# 
#     # Interpolate ps onto lat-lon grid
#     ps = interp_latlon(ps, lat, lon)
# 
#     return ps
# 
# def correct_for_topography(data, topo_ps, plev=None, lat=None, lon=None):
#     """Set pressure level data below topography to NaN.
# 
#     Parameters
#     ----------
#     data : ndarray or xray.DataArray
#         Data to correct, with pressure, latitude, longitude as the
#         last three dimensions.
#     topo_ps : ndarray or xray.DataArray
#         Climatological surface pressure to use for topography, on same
#         lat-lon grid as data.
#     plev, lat, lon : 1-D float array, optional
#         Pressure levels, latitudes and longitudes of input data.
#         Only used if data is an ndarray. If data is an xray.DataArray
#         then plev, lat and lon are extracted from data.coords.
# 
#     Returns
#     -------
#     data_out : ndarray or xray.DataArray
#         Data with grid points below topography set to NaN.
#     """
# 
#     if isinstance(data, xray.DataArray):
#         lat = get_coord(data, 'lat')
#         lon = get_coord(data, 'lon')
#         name, attrs, coords, _ = xr.meta(data)
#         vals = data.values.copy()
#         # -- Pressure levels in Pascals
#         plev = get_coord(data, 'plev')
#         pname = get_coord(data, 'plev', 'name')
#         plev = pres_convert(plev, data[pname].units, 'Pa')
#     else:
#         vals = data
# 
#     if isinstance(topo_ps, xray.DataArray):
#         if not latlon_equal(data, topo_ps):
#             msg = 'Inputs data and topo_ps are not on same latlon grid.'
#             raise ValueError(msg)
# 
#         # Surface pressure values in Pascals:
#         ps_vals = topo_ps.values
#         ps_vals = pres_convert(ps_vals, topo_ps.units, 'Pa')
#     else:
#         ps_vals = topo_ps
# 
#     # For each vertical level, set any point below topography to NaN
#     for k, p in enumerate(plev):
#         ibelow = ps_vals < p
#         vals[...,k,ibelow] = np.nan
# 
#     if isinstance(data, xray.DataArray):
#         data_out = xray.DataArray(vals, name=name, coords=coords, attrs=attrs)
#     else:
#         data_out = vals
# 
#     return data_out


import numpy as np
import matplotlib.pyplot as plt
import xray
import atmos.utils as utils
import atmos.plots as ap
import atmos.data as dat
from atmos.data import get_ps_clim, correct_for_topography


# ----------------------------------------------------------------------
# Read monthly mean climatologies and do some test calcs
# ----------------------------------------------------------------------

filename = 'data/more/ncep2_climatology_monthly.nc'
ds = dat.ncload(filename)
lat = ds['lat'].values
lon = ds['lon'].values
plev = ds['plev'].values
mon = ds['mon'].values
u = ds['u']
v = ds['v']
T = ds['T']

# Path: atmos/data.py
# def get_ps_clim(lat, lon, datafile='data/topo/ncep2_ps.nc'):
#     """Return surface pressure climatology on selected lat-lon grid.
# 
#     Parameters
#     ----------
#     lat, lon : 1-D float array
#         Latitude and longitude grid to interpolate surface pressure
#         climatology onto.
#     datafile : string, optional
#         Name of file to read for surface pressure climatology.
# 
#     Returns
#     -------
#     ps : xray.DataArray
#         DataArray of surface pressure climatology interpolated onto
#         lat-lon grid.
#     """
# 
#     ds = ncload(datafile)
#     ps = ds['ps']
#     ps.attrs = utils.odict_insert(ps.attrs, 'title', ds.attrs['title'], pos=0)
# 
#     # Check what longitude convention is used in the surface pressure
#     # climatology and switch if necessary
#     lonmax = lon_convention(lon)
#     lon_ps = get_coord(ps, 'lon')
#     if lon_convention(lon_ps) != lonmax:
#         ps = set_lon(ps, lonmax)
# 
#     # Interpolate ps onto lat-lon grid
#     ps = interp_latlon(ps, lat, lon)
# 
#     return ps
# 
# def correct_for_topography(data, topo_ps, plev=None, lat=None, lon=None):
#     """Set pressure level data below topography to NaN.
# 
#     Parameters
#     ----------
#     data : ndarray or xray.DataArray
#         Data to correct, with pressure, latitude, longitude as the
#         last three dimensions.
#     topo_ps : ndarray or xray.DataArray
#         Climatological surface pressure to use for topography, on same
#         lat-lon grid as data.
#     plev, lat, lon : 1-D float array, optional
#         Pressure levels, latitudes and longitudes of input data.
#         Only used if data is an ndarray. If data is an xray.DataArray
#         then plev, lat and lon are extracted from data.coords.
# 
#     Returns
#     -------
#     data_out : ndarray or xray.DataArray
#         Data with grid points below topography set to NaN.
#     """
# 
#     if isinstance(data, xray.DataArray):
#         lat = get_coord(data, 'lat')
#         lon = get_coord(data, 'lon')
#         name, attrs, coords, _ = xr.meta(data)
#         vals = data.values.copy()
#         # -- Pressure levels in Pascals
#         plev = get_coord(data, 'plev')
#         pname = get_coord(data, 'plev', 'name')
#         plev = pres_convert(plev, data[pname].units, 'Pa')
#     else:
#         vals = data
# 
#     if isinstance(topo_ps, xray.DataArray):
#         if not latlon_equal(data, topo_ps):
#             msg = 'Inputs data and topo_ps are not on same latlon grid.'
#             raise ValueError(msg)
# 
#         # Surface pressure values in Pascals:
#         ps_vals = topo_ps.values
#         ps_vals = pres_convert(ps_vals, topo_ps.units, 'Pa')
#     else:
#         ps_vals = topo_ps
# 
#     # For each vertical level, set any point below topography to NaN
#     for k, p in enumerate(plev):
#         ibelow = ps_vals < p
#         vals[...,k,ibelow] = np.nan
# 
#     if isinstance(data, xray.DataArray):
#         data_out = xray.DataArray(vals, name=name, coords=coords, attrs=attrs)
#     else:
#         data_out = vals
# 
#     return data_out


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



# ----------------------------------------------------------------------
# Read monthly mean climatologies and do some test calcs
# ----------------------------------------------------------------------

filename = 'data/more/ncep2_climatology_monthly.nc'
ds = dat.ncload(filename)
lat = ds['lat'].values
lon = ds['lon'].values
plev = ds['plev'].values
mon = ds['mon'].values
u = ds['u']
v = ds['v']
T = ds['T']

ps = ds['ps']