#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()
    _v = __prepare_gradient_arg(_fn_name, 1, _args)
    if count(_v) <> 4 then
        fail(_fn_name & ": invalid arguments=" & _args)
    end if
    _z = _v[1]
    _mode = _v[2]
    _pole_missing = _v[3]
    _vector_mode = _v[4]

    if count(_z) = 0 then
		fail(_fn_name & ": no valid input fieldset is specified")
	end if

    # some constants
	_omega = 7.292115E-5 # angular velocity of Earth
    _coriolis_limit = 1E-7
       
    # extract metadata keys
    _keys = grib_get(_z,["gridType","typeOfLevel"])
    
    # check inout 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 _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

    # compute the Coriolis parameter
    _f=(2*_omega)*sinlat(_z)

    # mask the tropics
    _f = bitmap(_f,bitmap(abs(_f) > _coriolis_limit,0))

    # compute the geopotential gradient
    if _mode = "fdiff" then
        _grad_z = gradient(_z)
    else
        _grad_z = regrid(data: _z, nabla: "scalar_gradient", nabla_poles_missing_values: _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
