#Metview Macro

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


# **************************************************************************
# Computes the static stability uused for the Q-vector computations3
#
# OneLineDesc   : Computes the Q-vector for temperature and geopotential fieldsets
#
# **************************************************************************

function q_vector
    _fn_name = "q_vector"
    
    _args = arguments()
    _v = __prepare_gradient_arg(_fn_name, 3, _args)
    if count(_v) <> 6 then
        fail(_fn_name & ": invalid arguments=" & _args)
    end if
    _t = _v[1]
    _z = _v[2]
    _sigma = _v[3]
    _mode = _v[4]
    _pole_missing = _v[5]
    _vector_mode = _v[6]

    gribsetbits(24)

    R = 8.314462 # 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 
    if count(_t) <> count(_sigma) then
        fail(_fn_name & ": number of temperature fields" & 
                    " (=" & count(_t) & 
                    ") is different than the number of static_stability fields (=" & count(_sigma) & ")!")
    end if 

    _p = __get_pressure_from_pl_arg(_t, "t", _fn_name)    

    _wg = geostrophic_wind(_z, "mode", _mode, "poles_missing_vales", _pole_missing, "vector", _vector_mode)
    _d_wg = gradient(_wg, "mode", _mode, "poles_missing_vales", _pole_missing, "vector", _vector_mode)
    _d_t = gradient(_t, "mode", _mode, "poles_missing_vales", _pole_missing, "vector", _vector_mode)
    
    _res = nil 
    for _i=1 to count(_t) do
        _c = (-R /_p[_i])/_sigma[_i]
        _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
