#Metview Macro

#  **************************** LICENSE START ***********************************
# 
#  Copyright 2022 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 ************************************
# 


# **************************************************************************
#
# OneLineDesc   : Computes the Q-vector for fieldsets
#
# **************************************************************************

function q_vector
    _fn_name = "q_vector"
    
    _args = arguments()
    
    if count(_args) < 2 then 
        fail(_fn_name & ": too few arguments specified!")
    end if
    
    _t = _args[1]
    __check_arg_type(_fn_name, _t, "fieldset", "temperature")
    
    _z = _args[2]
    __check_arg_type(_fn_name, _z, "fieldset", "geopotential")
    
    # 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 = (static_stability: 2E-6, coriolis: 1000, mode: "fdiff", poles_missing_values: "off", vector: "off")
    _opt = __get_kwargs(_fn_name, _args, 3, _opt)
   
    #print("opt=", _opt)
    # get static stability 
    _sigma = _opt.static_stability
    _opt.sigma = nil
    __check_arg_type(_fn_name, _sigma, ["number", "list", "vector", "fieldset"], 
                     "static stability")
    
    if type(_sigma) = "number" then
        _sigma = [_sigma]
    end if 
       
    # coriolis
    _coriolis = _opt.const_coriolis
   
    # the values can be very small so we increase the precision
    gribsetbits(24)

    Rd = 287.0597 # gas constant
    
    if count(_t) <> count(_z) then
        fail(_fn_name & ": number of temperature fields" & 
             " (=" & count(_t) & 
             ") is different than the number of geopotential fields (=" & count(_z) & ")!")
    end if
    
    _single_sigma = 0
    if count(_sigma) = 1 then
        _single_sigma = 1
    end if
    
    if _single_sigma = 0 and count(_t) <> count(_sigma) then
        fail(_fn_name & ": number of temperature fields" & 
                    " (=" & count(_t) & 
                    ") is different than the number of static_stability values/fields (=" & count(_sigma) & ")!")
    end if
   
    _p = __get_pressure_from_pl_arg(_t, "t", _fn_name)    

    _wg = geostrophic_wind(_z, 
                    "coriolis", _coriolis,
                    "mode", _opt.mode, 
                    "poles_missing_values", _opt.poles_missing_values, 
                    "vector", _opt.vector)
    _d_wg = gradient(_wg,  
                    "mode", _opt.mode, 
                    "poles_missing_values", _opt.poles_missing_values, 
                    "vector", _opt.vector)
    _d_t = gradient(_t,  
                    "mode", _opt.mode, 
                    "poles_missing_values", _opt.poles_missing_values, 
                    "vector", _opt.vector)
    
    _res = nil 
    for _i=1 to count(_t) do
        _sigma_idx = 1
        if _single_sigma = 0 then
            _sigma_idx = i
        end if
        _c = (-Rd /_p[_i])/_sigma[_sigma_idx]
        
        _t_idx_x = 2*_i-1
        _t_idx_y = 2*_i
        _u_idx_x = 4*_i-3
        _u_idx_y = 4*_i-2
        _v_idx_x = 4*_i-1
        _v_idx_y = 4*_i
        _q1 = (_d_wg[_u_idx_x] * _d_t[_t_idx_x] + _d_wg[_v_idx_x] * _d_t[_t_idx_y]) *_c
        # here we swap the order so that the metadata of v can be copied into q2
        _q2 = (_d_wg[_v_idx_y] * _d_t[_t_idx_y] + _d_wg[_u_idx_y] * _d_t[_t_idx_x]) *_c

        # create a unified bitmap
        _bm = _q1*0 + _q2*0
        _q1 = bitmap(_q1, _bm)
        _q2 = bitmap(_q2, _bm)

        _res = _res & _q1 & _q2
    end for
    
    return _res
end q_vector
