#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 ************************************
# 

# **************************************************************************
# Function      : geostrophic_wind_ml
#
# Syntax        : fieldset geostrophic_wind_ml(fs_x:fieldset,fs_y:fieldset)
#                                          
# Category      : DERIVATIVES
#
# OneLineDesc   : Compute geostrophic wind on (hybrid) model levels
#
# Description   : Compute the geostrophic wind on (hybrid) model levels
#                 on a 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_ml(z: fieldset, t: fieldset,q: fieldset, lnsp: fieldset)

	fn_name = "geostrophic_wind_ml"

    # some constants
	Rd = 8.3144598 # gas constant
	omega = 7.292115E-5 # angular velocity of Earth
    coriolis_limit = 1E-7
       
	# check fields          	 
	if count(z) = 0 then
		fail(fn_name & ": no valid input z fieldset is specifed")
	end if
         	             
    if count(t) <> count(z) then
        fail(fn_name & ": different number of fields in input fieldsets [z=",count(z),", t=",count(t),"]")
    end if	    	
            
    if count(q) <> count(z) then
        fail(fn_name & ": different number of fields in input fieldsets [z=",count(z),", q=",count(q),"]")
    end if 
          
    if count(lnsp) = 0 then
        fail(fn_name & ": no lnsp field is specified")
    end if   
  
    # extract metadata keys
    keys_z = grib_get(z,["gridType","typeOfLevel"])
    keys_t = grib_get(t,["gridType","typeOfLevel"])
    keys_q = grib_get(q,["gridType","typeOfLevel"])
    
    res = nil
    
    # compute the pressure on model levels
    p = unipressure(lnsp, t)
   
    for i=1 to count(z) do

    	# get metadata keys
        grid_type_z = keys_z[i][1]
        level_type_z = keys_z[i][2]
        grid_type_t = keys_t[i][1]
        level_type_t = keys_t[i][2]
        grid_type_q = keys_q[i][1]
        level_type_q = keys_q[i][2]
        
        # check if level is model level
        if level_type_z <> "hybrid" then
            fail(fn_name & ": [z field=",i,"] - unsupported level type(=",level_type_z,"), must be \"hybrid\" ")
        end if
        
        if level_type_t <> "hybrid" then
            fail(fn_name & ": [t field=",i,"] - unsupported level type(=",level_type_t,"), must be \"hybrid\" ")
        end if
        
        if level_type_q <> "hybrid" then
            fail(fn_name & ": [q field=",i,"] - unsupported level type(=",level_type_q,"), must be \"hybrid\" ")
        end if
        
        # check if grid is regular latlon 
        if grid_type_z <> "regular_ll" then
            fail(fn_name & ": [z field=",i,"] - unsupported grid (=",grid_type_z,"), implemented only for regular lat-lon grid")
        end if
                
        if grid_type_t <> "regular_ll" then
            fail(fn_name & ": [t field=",i,"] - unsupported grid (=",grid_type_t,"), implemented only for regular lat-lon grid")
        end if
  
        if grid_type_q <> "regular_ll" then
            fail(fn_name & ": [q field=",i,"] - unsupported grid (=",grid_type_q,"), implemented only for regular lat-lon grid")
        end if
        
        # compute the Coriolis parameter
        f=2*omega*sinlat(z[i])

        # mask the tropics
        f = bitmap(f,bitmap(abs(f) > coriolis_limit,0))

        # compute the virtual temperature
    	tv = t[i] * (1  + 0.61*q[i])

        # compute the geopotential gradient
        grad_z = gradient(z[i])

        # compute the pressure gradient
        grad_p = gradient(log(p[i]))

        # compute the geostrophic wind
        u_g = -grad_z[2]/f - grad_p[2] * Rd * tv / f 
        v_g = grad_z[1]/f + grad_p[1] * Rd * tv / f
	
		# set param ids as U/V according to level    
        u_g = grib_set_long(u_g,["paramId",131])
        v_g = grib_set_long(v_g,["paramId",132])
       
        res = res & u_g & v_g
		
	end for
	
    return res
    
end geostrophic_wind_ml
