from __future__ import (absolute_import, division, print_function, 
                        unicode_literals)

from .constants import Constants
from .destag import destagger
from .decorators import convert_units
from .metadecorators import set_height_metadata
from .util import extract_vars, either

def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, 
               cache=None, meta=True, _key=None,
               height=True, msl=True):
    """Return the geopotential or geopotential height.
    
    If *height* is False, then geopotential is returned in units of 
    [m2 s-2].  If *height* is True, then geopotential height is 
    returned in units of [m].  If *msl* is True, then geopotential height 
    is return as Mean Sea Level (MSL).  If *msl* is False, then geopotential 
    height is returned as Above Ground Level (AGL).
    
    This functions extracts the necessary variables from the NetCDF file 
    object in order to perform the calculation.
    
    Args:
    
        wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
            iterable): WRF-ARW NetCDF 
            data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` 
            or an iterable sequence of the aforementioned types.
        
        timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The 
            desired time index. This value can be a positive integer, 
            negative integer, or 
            :data:`wrf.ALL_TIMES` (an alias for None) to return 
            all times in the file or sequence. The default is 0.
        
        method (:obj:`str`, optional): The aggregation method to use for 
            sequences.  Must be either 'cat' or 'join'.  
            'cat' combines the data along the Time dimension.  
            'join' creates a new dimension for the file index.  
            The default is 'cat'.
        
        squeeze (:obj:`bool`, optional): Set to False to prevent dimensions 
            with a size of 1 from being automatically removed from the shape 
            of the output. Default is True.
        
        cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) 
            that can be used to supply pre-extracted NetCDF variables to the 
            computational routines.  It is primarily used for internal 
            purposes, but can also be used to improve performance by 
            eliminating the need to repeatedly extract the same variables 
            used in multiple diagnostics calculations, particularly when using 
            large sequences of files. 
            Default is None.
        
        meta (:obj:`bool`, optional): Set to False to disable metadata and 
            return :class:`numpy.ndarray` instead of 
            :class:`xarray.DataArray`.  Default is True.
            
        _key (:obj:`int`, optional): A caching key. This is used for internal 
            purposes only.  Default is None.
            
        height (:obj:`bool`, optional): Set to True to return geopotential 
            height instead of geopotential.  Default is True.
            
        msl (:obj:`bool`, optional): Set to True to return geopotential height
            as Mean Sea Level (MSL).  Set to False to return the 
            geopotential height as Above Ground Level (AGL) by subtracting 
            the terrain height.  Default is True.
   
    Returns:
        :class:`xarray.DataArray` or :class:`numpy.ndarray`: The 
        geopotential or geopotential height.
        If xarray is enabled and the *meta* parameter is True, then the result 
        will be a :class:`xarray.DataArray` object.  Otherwise, the result will 
        be a :class:`numpy.ndarray` object with no metadata.
    
    """
    
    varname = either("PH", "GHT")(wrfin)
    if varname == "PH":
        ph_vars = extract_vars(wrfin, timeidx, ("PH", "PHB", "HGT"),
                               method, squeeze, cache, meta=False,
                               _key=_key)
        ph = ph_vars["PH"]
        phb = ph_vars["PHB"]
        hgt = ph_vars["HGT"]
        geopt = ph + phb
        geopt_unstag = destagger(geopt, -3)
    else:
        ght_vars = extract_vars(wrfin, timeidx, ("GHT", "HGT_M"),
                                method, squeeze, cache, meta=False,
                                _key=_key)
        geopt_unstag = ght_vars["GHT"] * Constants.G
        hgt = ght_vars["HGT_M"]
    
    if height:
        if msl:
            return geopt_unstag / Constants.G
        else:
            # Due to broadcasting with multifile/multitime, the 2D terrain
            # array needs to be reshaped to a 3D array so the right dims
            # line up
            new_dims = list(hgt.shape)
            new_dims.insert(-2,1)
            hgt = hgt.reshape(new_dims)
            
            return (geopt_unstag / Constants.G) - hgt
    else:
        return geopt_unstag


@set_height_metadata(geopt=True)
def get_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, 
              meta=True, _key=None):
    """Return the geopotential.
    
    The geopotential is returned in units of [m2 s-2].
    
    This functions extracts the necessary variables from the NetCDF file 
    object in order to perform the calculation.
    
    Args:
    
        wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
            iterable): WRF-ARW NetCDF 
            data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` 
            or an iterable sequence of the aforementioned types.
        
        timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The 
            desired time index. This value can be a positive integer, 
            negative integer, or 
            :data:`wrf.ALL_TIMES` (an alias for None) to return 
            all times in the file or sequence. The default is 0.
        
        method (:obj:`str`, optional): The aggregation method to use for 
            sequences.  Must be either 'cat' or 'join'.  
            'cat' combines the data along the Time dimension.  
            'join' creates a new dimension for the file index.  
            The default is 'cat'.
        
        squeeze (:obj:`bool`, optional): Set to False to prevent dimensions 
            with a size of 1 from being automatically removed from the shape 
            of the output. Default is True.
        
        cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) 
            that can be used to supply pre-extracted NetCDF variables to the 
            computational routines.  It is primarily used for internal 
            purposes, but can also be used to improve performance by 
            eliminating the need to repeatedly extract the same variables 
            used in multiple diagnostics calculations, particularly when using 
            large sequences of files. 
            Default is None.
        
        meta (:obj:`bool`, optional): Set to False to disable metadata and 
            return :class:`numpy.ndarray` instead of 
            :class:`xarray.DataArray`.  Default is True.
            
        _key (:obj:`int`, optional): A caching key. This is used for internal 
            purposes only.  Default is None.
   
    Returns:
        :class:`xarray.DataArray` or :class:`numpy.ndarray`: The 
        geopotential.
        If xarray is enabled and the *meta* parameter is True, then the result 
        will be a :class:`xarray.DataArray` object.  Otherwise, the result will 
        be a :class:`numpy.ndarray` object with no metadata.
    
    """
    return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key,
                      False, True)


@set_height_metadata(geopt=False)
@convert_units("height", "m")
def get_height(wrfin, timeidx=0, method="cat", squeeze=True, 
               cache=None, meta=True, _key=None,
               msl=True, units="m"):
    """Return the geopotential height.
    
    If *msl* is True, then geopotential height is returned as Mean Sea Level 
    (MSL).  If *msl* is False, then geopotential height is returned as 
    Above Ground Level (AGL) by subtracting the terrain height.
    
    This functions extracts the necessary variables from the NetCDF file 
    object in order to perform the calculation.
    
    Args:
    
        wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
            iterable): WRF-ARW NetCDF 
            data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` 
            or an iterable sequence of the aforementioned types.
        
        timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The 
            desired time index. This value can be a positive integer, 
            negative integer, or 
            :data:`wrf.ALL_TIMES` (an alias for None) to return 
            all times in the file or sequence. The default is 0.
        
        method (:obj:`str`, optional): The aggregation method to use for 
            sequences.  Must be either 'cat' or 'join'.  
            'cat' combines the data along the Time dimension.  
            'join' creates a new dimension for the file index.  
            The default is 'cat'.
        
        squeeze (:obj:`bool`, optional): Set to False to prevent dimensions 
            with a size of 1 from being automatically removed from the shape 
            of the output. Default is True.
        
        cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) 
            that can be used to supply pre-extracted NetCDF variables to the 
            computational routines.  It is primarily used for internal 
            purposes, but can also be used to improve performance by 
            eliminating the need to repeatedly extract the same variables 
            used in multiple diagnostics calculations, particularly when using 
            large sequences of files. 
            Default is None.
        
        meta (:obj:`bool`, optional): Set to False to disable metadata and 
            return :class:`numpy.ndarray` instead of 
            :class:`xarray.DataArray`.  Default is True.
            
        _key (:obj:`int`, optional): A caching key. This is used for internal 
            purposes only.  Default is None.
            
        msl (:obj:`bool`, optional): Set to True to return geopotential height
            as Mean Sea Level (MSL).  Set to False to return the 
            geopotential height as Above Ground Level (AGL) by subtracting 
            the terrain height.  Default is True.
           
        units (:obj:`str`): The desired units.  Refer to the :meth:`getvar` 
            product table for a list of available units for 'z'.  Default 
            is 'm'.
   
    Returns:
        :class:`xarray.DataArray` or :class:`numpy.ndarray`: The 
        geopotential height.
        If xarray is enabled and the *meta* parameter is True, then the result 
        will be a :class:`xarray.DataArray` object.  Otherwise, the result will 
        be a :class:`numpy.ndarray` object with no metadata.
    
    """
    
    return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key,
                      True, msl)