#Metview Macro

#  **************************** LICENSE START ***********************************
# 
#  Copyright 2018 ECMWF. This software is distributed under the terms
#  of the Apache License version 2.0. In applying this license, ECMWF does not
#  waive the privileges and immunities granted to it by virtue of its status as
#  an Intergovernmental Organization or submit itself to any jurisdiction.
# 
#  ***************************** LICENSE END ************************************
# 

# **************************************************************************
# Function      : geostrophic_wind
#
# Syntax        : fieldset geostrophic_wind(fs_x:fieldset,fs_y:fieldset)
#                                          
# Category      : DERIVATIVES
#
# OneLineDesc   : Computes geostrophic wind on pressure levels
#
# Description   : Compute the geostrophic wind from geopotentials fields defined on 
#                 pressure levels and regular lat-lon grid. The computation uses a 
#                 second order accuracy
#                 finite difference scheme.
#
# Parameters    : 
#                 
# Return Value  : resulting fieldset
#
# Dependencies  : none
#
# Example Usage : 
#                 
#
# **************************************************************************
function geostrophic_wind
    _fn_name = "geostrophic_wind"
    
    _args = arguments()
     
    _z = _args[1]
    __check_arg_type(_fn_name, _z, "fieldset", "geopotential")
    
    if count(_z) = 0 then
		fail(_fn_name & ": no valid input fieldset is specified")
	end if

    # keyword args
    # coriolis default is nil. However, if we set the value of a key in a definition
    # to nil the key will be removed! So we use a fill value here.
    _opt = (coriolis: 1000, mode: "fdiff", poles_missing_values: "off", vector: "off")
    _opt = __get_kwargs(_fn_name, _args, 2, _opt)

    # some constants
	_omega = 7.29211508E-5 # angular velocity of Earth
    _coriolis_limit = 1E-7
       
    # extract metadata keys
    _keys = grib_get(_z,["gridType","typeOfLevel"])
    
    # check input data
    for _i=1 to count(_z) do
    	# get metadata keys
        _grid_type = _keys[_i][1]
        _level_type = _keys[_i][2]
        
        # check if level is pressure 
        if _level_type <> "isobaricInhPa" then
            fail(_fn_name & ": [z field=",_i,"] - unsupported level type(=",_level_type,"), must be \"isobaricInhPa\" ")
        end if
        
        # check if grid is regular latlon 
        if _opt.mode = "fdiff" then
            if _grid_type <> "regular_ll" then
                fail(_fn_name & ": [z field=",_i,"] - unsupported grid (=",_grid_type,"), implemented only for regular lat-lon grid")
            end if
        end if
    end for
   
    _corilolis_is_number = 0
    if type(_opt.coriolis) = "number" then
        _f = _opt.coriolis
        if abs(_f-1000) > 1E-4 then
            _corilolis_is_number = 1
            if abs(_f) > 2*_omega then
                fail(_fn_name &": invalid coriolis value=" & 
                    "Valid interval=[" & -2*_omega & "," & 2*_omega & "]")
            end if
        end if
    end if     
    if _corilolis_is_number = 0 then   
        # compute the Coriolis parameter
        _f=(2*_omega)*sinlat(_z)

        # mask the tropics
        _f = bitmap(_f,bitmap(abs(_f) > _coriolis_limit,0))
    end if
    
    # compute the geopotential gradient
    if _opt.mode = "fdiff" then
        _grad_z = gradient(_z)
    else
        _grad_z = regrid(data: _z, nabla: "scalar_gradient", nabla_poles_missing_values: opt._pole_missing)
    end if

    # compute the geostrophic wind
    _u_g = -_grad_z[2, count(_grad_z), 2]/_f
    _v_g = _grad_z[1, count(_grad_z)-1,2]/_f

    # create a unified bitmap
    _bm = _u_g*0 + _v_g*0
    _u_g = bitmap(_u_g, _bm)
    _v_g = bitmap(_v_g, _bm)

    # set param ids as U/V  
    _u_g = grib_set_long(_u_g,["paramId",131])
    _v_g = grib_set_long(_v_g,["paramId",132])
      
    _res = nil
    for _i = 1 to count(_u_g) do
        _res = _res & _u_g[_i] & _v_g[_i]	
	end for
	
    return _res
    
end geostrophic_wind
