#Metview Macro

#  **************************** LICENSE START ***********************************
# 
#  Copyright 2019 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 vertical pressure derivative
#
# OneLineDesc   : Computes the vertical pressure derivative
#
# Input: 
#       fs: fieldset
#       p: the pressure (Pa)
# Return:
#       the pressure derivative        
#==============================================================================

function pressure_derivative(fs)
    fn_name = "pressure_derivative"   
    p = __get_pressure_from_pl_arg(fs, "fs", fn_name)       
    return pressure_derivative(fs, p)  
      
end pressure_derivative

function pressure_derivative(fs, p)

    __DEBUG_PRESSURE_DERIVATIVE = 0

    fn_name = "pressure_derivative"   
    has_pl_fields = 0
    dp_eps = 1E-3
       
    if  __DEBUG_PRESSURE_DERIVATIVE then
        print("p type=", type(p))
    end if

    if count(p) < 2 then
        fail(fn_name  & ":number of levels (=" & count(p) & ") must be >=2")
    end if

    if type(fs) = "fieldset" then
        v = __prepare_pressure_field_arg(fs, p, "fs", fn_name)
        p = v[1]
        has_pl_fields = v[2]
    end if     
    
    # we need to use it because when called from Python
    # sort_indices() returns 0-based index values and it cannot be used
    # for indexing fieldsets, which expects 1-based values.
    # So we need to adjust the result of sort_indices() to make
    # pressure_derivative() work in Python.
    fs_index_offset = 0
    if base_language = 'python' then 
	    fs_index_offset = -1
	end if

    # Sort by pressure (descending)
    # We do not sort the data itself to avoid I/O 
    # but store the sort indices!
    if has_pl_fields or type(p) <> "fieldset" then
        fIndex = sort_indices(p, ">") 
    else if type(p) = "fieldset" then
        p_avg = average(p)
        fIndex = sort_indices(p_avg, ">") 
    end if

    # we need to adjust these values to make the function work from Python
    fIndex = fIndex - fs_index_offset 
    
    # generate the pIndex
    pIndex = fIndex
    for i=1 to count(fIndex) do
        pIndex[fIndex[i]] =i
    end for

    if  __DEBUG_PRESSURE_DERIVATIVE then
        print("pIndex=", pIndex)
        print("fIndex=", fIndex)
    end if

    if type(fs) = "fieldset" then
        r = nil
    else
        r =  fs*0
    end if

    num = count(fIndex)
    if num <> count(fs) then
        fail(fn_name & ": inconsistent number level of indices " & num  & "!=" & count(fs))
    end if

    # fIndex[i]: at the i-th sorted pressure gives the position of the 
    #            corresponding field (fs[fIndex[i]])
    # pIndex[i]: at the i-th field position (fs[i]) gives the sorted pressure position

    # Compute the derivative. We need to keep the original field ordering in the output!
    for i=1 to count(pIndex) do 
        idx = pIndex[i]
        
        if  __DEBUG_PRESSURE_DERIVATIVE then
            print("i=", i, " idx=", idx)
        end if

        # The first field in the equation has to contain the derivative grib
        # metadata for that level!

        # bottom - one sided difference
        if idx = 1 then       
            idx0 = fIndex[1]
            idx1 = fIndex[2]
            dp =  p[idx0] - p[idx1]
            
            if  __DEBUG_PRESSURE_DERIVATIVE then
                print("  p0=", p[idx0], " p1=",  p[idx1])
            end if

            if type(p) <> "fieldset" then 
                if abs(dp) < dp_eps then
                    fail(fn_name & ": identical pressure levels found (p=" & p[idx0] & "Pa)!")
                end if
            else
                dp = dp*bitmap(abs(dp) > dp_eps, 0)
            end if
            if type(fs) = "fieldset" then
                r = r & ((fs[idx0] - fs[idx1]) / dp)
            else
                r[i] = (fs[idx0] - fs[idx1]) / dp
            end if
                 
        # mid levels - central difference
        else if idx < count(pIndex) then
            idx0 = fIndex[idx-1]
            idx1 = fIndex[idx]
            idx2 = fIndex[idx+1]
            dp = p[idx2] - p[idx0]
            if  __DEBUG_PRESSURE_DERIVATIVE then
                print("  p1=", p[idx1], "  p0=",  p[idx0], " p2=", p[idx2])
            end if
            if abs(dp) < 1E-3 then
                fail(fn_name & ": identical pressure levels found (p=" & p[idx0] & "Pa)!")
            end if
            if type(fs) = "fieldset" then
                r = r & (fs[idx1]*0 +  (fs[idx2] - fs[idx0]) / dp)
            else
                r[i] = (fs[idx2] - fs[idx0]) / dp
            end if
        
        # top - one sided difference
        else if idx = num then
            idx0 = fIndex[idx-1]
            idx1 = fIndex[idx]
            dp = p[idx1] - p[idx0]
            if  __DEBUG_PRESSURE_DERIVATIVE then
                print("  p1=", p[idx1], " p0=", p[idx0])
            end if
            if abs(dp) < 1E-3 then
                fail(fn_name & ": identical pressure levels found (p=" & p[idx0] & "Pa)!")
            end if
            if type(fs) = "fieldset" then
                r = r & ((fs[idx1] - fs[idx0]) / dp)
            else
                r[i] = (fs[idx1] - fs[idx0]) / dp
            end if
        else
            fail(fn_name & ": invalid index in computations " & idx & "=" & idx)       
        end if
    end for
    
    return r

end pressure_derivative
