#include "../includeF/sfx_def.fi"

module sfx_most_snow
    !< @brief MOST surface flux module

    ! modules used
    ! --------------------------------------------------------------------------------
#ifdef SFX_CHECK_NAN
    use sfx_common
#endif
    use sfx_data
    use sfx_surface
    use sfx_most_snow_param
    use sfx_config
    ! --------------------------------------------------------------------------------

    ! directives list
    ! --------------------------------------------------------------------------------
    implicit none
    private
    ! --------------------------------------------------------------------------------

    ! public interface
    ! --------------------------------------------------------------------------------
    public :: get_surface_fluxes
    public :: get_surface_fluxes_vec
    public :: get_psi
    integer z0m_id
    integer z0t_id
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    type, public :: numericsType
        integer :: maxiters_charnock = 10      !< maximum (actual) number of iterations in charnock roughness
        integer :: maxiters_snow = 10          !< maximum (actual) number of iterations in snow roughness
    end type
    ! --------------------------------------------------------------------------------

contains

    ! --------------------------------------------------------------------------------
    subroutine get_surface_fluxes_vec(sfx, meteo, numerics, n)
        !< @brief surface flux calculation for array data
        !< @details contains C/C++ & CUDA interface
        ! ----------------------------------------------------------------------------
        type (sfxDataVecType), intent(inout) :: sfx

        type (meteoDataVecType), intent(in) :: meteo
        type (numericsType), intent(in) :: numerics
        integer, intent(in) :: n
        ! ----------------------------------------------------------------------------

        ! --- local variables
        type (meteoDataType)  meteo_cell
        type (sfxDataType) sfx_cell
	
        integer i
        ! ----------------------------------------------------------------------------

        do i = 1, n
            meteo_cell = meteoDataType(&
                    h = meteo%h(i), &
                    U = meteo%U(i), dT = meteo%dT(i), Tsemi = meteo%Tsemi(i), dQ = meteo%dQ(i), &
                    z0_m = meteo%z0_m(i), depth=meteo%depth(i), lai=meteo%lai(i), surface_type=meteo%surface_type(i))

            call get_surface_fluxes(sfx_cell, meteo_cell, numerics)

            call push_sfx_data(sfx, sfx_cell, i)
            
		end do

    end subroutine get_surface_fluxes_vec
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    subroutine get_surface_fluxes(sfx, meteo, numerics)
        !< @brief surface flux calculation for single cell
        !< @details contains C/C++ interface
        ! ----------------------------------------------------------------------------
#ifdef SFX_CHECK_NAN
        use ieee_arithmetic
#endif

        type (sfxDataType), intent(out) :: sfx
        type (meteoDataType), intent(in) :: meteo
        type (numericsType), intent(in) :: numerics
        ! ----------------------------------------------------------------------------

        ! --- meteo derived datatype name shadowing
        ! ----------------------------------------------------------------------------
        real :: h       !< constant flux layer height [m]
        real :: U       !< abs(wind speed) at 'h' [m/s]
        real :: dT      !< difference between potential temperature at 'h' and at surface [K]
        real :: Tsemi   !< semi-sum of potential temperature at 'h' and at surface [K]
        real :: dQ      !< difference between humidity at 'h' and at surface [g/g]
        real :: z0_m    !< surface aerodynamic roughness (should be < 0 for water bodies surface)
        real :: depth  
        real :: lai
        ! ----------------------------------------------------------------------------

        ! --- local variables
        ! ----------------------------------------------------------------------------
        real z0_t               !< thermal roughness [m]
        real B                  !< = ln(z0_m / z0_t) [n/d]
        real h0_m, h0_t         !< = h / z0_m, h / z0_h [n/d]

        real u_dyn0             !< dynamic velocity in neutral conditions [m/s]
        real Re                 !< roughness Reynolds number = u_dyn0 * z0_m / nu [n/d]

        real zeta               !< = z/L [n/d]
        real Rib                !< bulk Richardson number

        real Udyn, Tdyn, Qdyn   !< dynamic scales

        real phi_m, phi_h       !< stability functions (momentum) & (heat) [n/d]

        real Km                 !< eddy viscosity coeff. at h [m^2/s]
        real Pr_t_inv           !< invese Prandt number [n/d]

        real Cm, Ct             !< transfer coeff. for (momentum) & (heat) [n/d]

        ! --- local additional variables
        ! ----------------------------------------------------------------------------
		!real phi_m, phi_m
		real hfx, mfx
		real S_mean, Lsnow
	    real z0_s, h_salt
                
	   
		integer surface_type    !< surface type = (ocean || land)


#ifdef SFX_CHECK_NAN
        real NaN
#endif
        ! ----------------------------------------------------------------------------

#ifdef SFX_CHECK_NAN
        ! --- checking if arguments are finite
        if (.not.(is_finite(meteo%U).and.is_finite(meteo%Tsemi).and.is_finite(meteo%dT).and.is_finite(meteo%dQ) &
                .and.is_finite(meteo%z0_m).and.is_finite(meteo%h))) then

            !NaN = ieee_value(0.0, ieee_quiet_nan)   ! setting NaN
            sfx = sfxDataType(zeta = NaN, Rib = NaN, &
                    Re = NaN, B = NaN, z0_m = NaN, z0_t = NaN, &
                    Rib_conv_lim = NaN, &
                    Cm = NaN, Ct = NaN, Km = NaN, Pr_t_inv = NaN)
            
            
            
            return
        end if
#endif

        ! --- shadowing names for clarity
        U = meteo%U
        Tsemi = meteo%Tsemi
        dT = meteo%dT
        dQ = meteo%dQ
        h = meteo%h
        z0_m = meteo%z0_m
        depth = meteo%depth
        lai = meteo%lai
        surface_type=meteo%surface_type

        
        call get_dynamic_roughness_definition(surface_type, ocean_z0m_id, land_z0m_id, lake_z0m_id, snow_z0m_id, &
        forest_z0m_id, usersf_z0m_id, z0m_id)
        
        call get_dynamic_roughness_all(z0_m, u_dyn0, U, depth, h, numerics%maxiters_charnock, z0_m, z0m_id)

        call get_thermal_roughness_definition(surface_type, ocean_z0t_id, land_z0t_id, lake_z0t_id, snow_z0t_id, &
        forest_z0t_id, usersf_z0t_id, z0t_id)
        
        Re = u_dyn0 * z0_m / nu_air
        
        call get_thermal_roughness_all(z0_t, B, z0_m, Re, u_dyn0, lai, z0t_id)
       
       ! --- define relative height
            h0_m = h / z0_m
        ! --- define relative height [thermal]
        h0_t = h / z0_t

        ! --- define Ri-bulk
        Rib = (g / Tsemi) * h * (dT + 0.61e0 * Tsemi * dQ) / U**2
        
		
		! --- get the fluxes
        ! ----------------------------------------------------------------------------
        call get_dynamic_scales(Udyn, Tdyn, Qdyn, zeta, &
	        Lsnow, S_mean, h_salt, &
            U, Tsemi, dT, dQ, h, z0_m, z0_t, (g / Tsemi), 10)
        ! ----------------------------------------------------------------------------

        call get_phi(phi_m, phi_h, zeta)
        ! ----------------------------------------------------------------------------

        ! --- define transfer coeff. (momentum) & (heat)
        Cm = 0.0
        if (U > 0.0) then
            Cm = Udyn / U
        end if
        Ct = 0.0
        if (abs(dT) > 0.0) then
            Ct = Tdyn / dT
        end if

        ! --- define eddy viscosity & inverse Prandtl number
        Km = kappa * Cm * U * h / phi_m
        Pr_t_inv = phi_m / phi_h
	    ! --- define heat flux and  momentum flux
		hfx=-rho_air*U*dT*Cm*Ct
		mfx=-rho_air*Cm*Cm*U*U
        h_salt=h_s

		! --- setting output
        sfx = sfxDataType(zeta = zeta, Rib = Rib, &
                Re = Re, B = B, z0_m = z0_m, z0_t = z0_t, &
                Rib_conv_lim = 0.0, &
                Cm = Cm, Ct = Ct, Km = Km, Pr_t_inv = Pr_t_inv)
		 ! --- setting additional output
		
				
    end subroutine get_surface_fluxes
    ! --------------------------------------------------------------------------------

    !< @brief get dynamic scales
    ! --------------------------------------------------------------------------------
    subroutine get_dynamic_scales(Udyn, Tdyn, Qdyn, zeta, &
	        Lsnow, S_mean, h_salt, &
            U, Tsemi, dT, dQ, z, z0_m, z0_t, beta, maxiters)
        ! ----------------------------------------------------------------------------
        real, intent(out) :: Udyn, Tdyn, Qdyn   !< dynamic scales
        real, intent(out) :: zeta               !< = z/L
		real, intent(out) :: Lsnow, S_mean             
        real, intent(out) :: h_salt        
		
        real, intent(in) :: U                   !< abs(wind speed) at z
        real, intent(in) :: Tsemi               !< semi-sum of temperature at z and at surface
        real, intent(in) :: dT, dQ              !< temperature & humidity difference between z and at surface
        real, intent(in) :: z                   !< constant flux layer height
        real, intent(in) :: z0_m, z0_t          !< roughness parameters
        real, intent(in) :: beta                !< buoyancy parameter

        integer, intent(in) :: maxiters         !< maximum number of iterations
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real, parameter :: gamma = 0.61

        real :: psi_m, psi_h
        real :: psi0_m, psi0_h
        real :: Linv
        real :: w_snow, sigma_m
        integer :: i
        ! ----------------------------------------------------------------------------


        Udyn = kappa * U / log(z / z0_m)
        Tdyn = kappa * dT * Pr_t_0_inv / log(z / z0_t)
        Qdyn = kappa * dQ * Pr_t_0_inv / log(z / z0_t)
        zeta = 0.0

        ! --- no wind
        if (Udyn < 1e-5) return

        Linv = kappa * beta * (Tdyn + gamma * Qdyn * Tsemi) / (Udyn * Udyn)
        zeta = z * Linv

        ! --- near neutral case
        if (Linv < 1e-5) return

        do i = 1, maxiters

            call get_psi(psi_m, psi_h, zeta)
            call get_psi_mh(psi0_m, psi0_h, z0_m * Linv, z0_t * Linv)

            Udyn = kappa * U / (log(z / z0_m) - (psi_m - psi0_m))
            Tdyn = kappa * dT * Pr_t_0_inv / (log(z / z0_t) - (psi_h - psi0_h))
            Qdyn = kappa * dQ * Pr_t_0_inv / (log(z / z0_t) - (psi_h - psi0_h))

            if (Udyn < 1e-5) exit

            Linv = kappa * beta * (Tdyn + gamma * Qdyn * Tsemi) / (Udyn * Udyn)
            zeta = z * Linv
       
            if (Udyn>u_thsnow) then
           
                call get_S_mean(S_mean, S_salt, h_s, z)
                call get_sigma_m(sigma_m, rho_air, rho_s)
                call get_w_snow(w_snow, sigma_m, g, d_s, nu_air)
                Linv=Linv*((1-S_mean)/(1+sigma_m*S_mean))+(g*w_snow*sigma_m*S_mean/(Udyn**3.0))/(1+sigma_m*S_mean)
                zeta = z * Linv
                Lsnow=1/Linv
				  !write(*,*) S_mean, sigma_m, w_snow
			      !pause 
			      !stop
				
				endif
       
       
       
        end do

    end subroutine get_dynamic_scales
    ! --------------------------------------------------------------------------------

    ! stability functions
    ! --------------------------------------------------------------------------------
    subroutine get_phi(phi_m, phi_h, zeta)
        !< @brief stability functions (momentum) & (heat): neutral case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: phi_m, phi_h   !< stability functions

        real, intent(in) :: zeta            !< = z/L
        ! ----------------------------------------------------------------------------


        if (zeta >= 0.0) then
            phi_m = 1.0 + beta_m * zeta
            phi_h = 1.0 + beta_h * zeta
        else
            phi_m = (1.0 - alpha_m * zeta)**(-0.25)
            phi_h = (1.0 - alpha_h * zeta)**(-0.5)
        end if

    end subroutine
    ! --------------------------------------------------------------------------------
    subroutine get_S_mean(S_mean, S_salt, h_s, z)
        !< @brief function for snow consentration
        ! ----------------------------------------------------------------------------
        real, intent(out) :: S_mean   !< snow consentration
        real, intent(in) ::  S_salt, h_s, z        
        ! ----------------------------------------------------------------------------
        S_mean  = S_salt *  h_s/z
    end subroutine
    ! --------------
    subroutine get_sigma_m(sigma_m, rho_air, rho_s)
        !< @brief function for 
        ! ----------------------------------------------------------------------------
        real, intent(out) :: sigma_m   !< s
        real, intent(in) ::  rho_air, rho_s            
        ! ----------------------------------------------------------------------------
        sigma_m  = (rho_s - rho_air)/rho_air
    end subroutine
    subroutine get_w_snow(w_snow, sigma_m, g, d_s, nu_air)
        !< @brief function for smow velosity
        ! ----------------------------------------------------------------------------
        real, intent(out) :: w_snow   !< 
        real, intent(in) ::  sigma_m, g, d_s, nu_air            
        ! ----------------------------------------------------------------------------
        w_snow  = sigma_m * g * d_s * d_s / (18.0 * nu_air);

    end subroutine
    ! universal functions
    ! --------------------------------------------------------------------------------
    subroutine get_psi(psi_m, psi_h, zeta)
        !< @brief universal functions (momentum) & (heat): neutral case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: psi_m, psi_h   !< universal functions

        real, intent(in) :: zeta            !< = z/L
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real :: x_m, x_h
        ! ----------------------------------------------------------------------------


        if (zeta >= 0.0) then
            psi_m = -beta_m * zeta;
            psi_h = -beta_h * zeta;
        else
            x_m = (1.0 - alpha_m * zeta)**(0.25)
            x_h = (1.0 - alpha_h * zeta)**(0.25)

            psi_m = (4.0 * atan(1.0) / 2.0) + 2.0 * log(0.5 * (1.0 + x_m)) + log(0.5 * (1.0 + x_m * x_m)) - 2.0 * atan(x_m)
            psi_h = 2.0 * log(0.5 * (1.0 + x_h * x_h))
        end if

    end subroutine


    subroutine get_psi_mh(psi_m, psi_h, zeta_m, zeta_h)
        !< @brief universal functions (momentum) & (heat): neutral case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: psi_m, psi_h   !< universal functions

        real, intent(in) :: zeta_m, zeta_h  !< = z/L
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real :: x_m, x_h
        ! ----------------------------------------------------------------------------


        if (zeta_m >= 0.0) then
            psi_m = -beta_m * zeta_m
        else
            x_m = (1.0 - alpha_m * zeta_m)**(0.25)
            psi_m = (4.0 * atan(1.0) / 2.0) + 2.0 * log(0.5 * (1.0 + x_m)) + log(0.5 * (1.0 + x_m * x_m)) - 2.0 * atan(x_m)
        end if

        if (zeta_h >= 0.0) then
            psi_h = -beta_h * zeta_h;
        else
            x_h = (1.0 - alpha_h * zeta_h)**(0.25)
            psi_h = 2.0 * log(0.5 * (1.0 + x_h * x_h))
        end if

    end subroutine
    ! --------------------------------------------------------------------------------

end module sfx_most_snow