!pbl_dry_contrgradient.f90
!>@brief implements mass-flux contr-gradient parametrization
module pbl_dry_contrgradient
    implicit none
    type pblContrGradDataType
        real, allocatable :: wu(:), &
                theta_up(:),  &
                thetav_up(:),  &
                ua_up(:), &
                va_up(:), &
                qv_up(:), &
                entr(:)
        real :: theta_up0, thetav_up0, qv_up0
        real :: ua_up0, va_up0
        real :: ua_upH, va_upH
        real :: theta_upH, thetav_upH, qv_upH
    end type pblContrGradDataType

    real, parameter :: c_entr = 0.38 !< dry entrainment parameterization
    real, parameter :: b1 = 1.8; !< vertical velocity in updraft Bernulli equation constant
    real, parameter :: b2 = 3.5; !< vertical velocity in updraft Bernulli equation constant
    real, parameter :: wstrmin = 1.0e-6 !< minimimal deardorff velocity
    real, parameter :: lmin_entr = 1.0e6 !< maximum entrainement scale above CBL
    real, parameter :: tau_dry_scale = 500.0 !< Time scale for entrainment
    real, parameter :: a0 = 0.08 !< Updraft fraction area
    real, parameter :: alph0 = 0.25 !< amount of surface dispersion to transfer to updraft

    public cbl_allocate, cbl_deallocate

    public cbl_get_entrainment, cbl_get_up
    public cbl_get_wu, cbl_apply_cntrg

    contains
        subroutine cbl_allocate(cbl, kmax)
            integer, intent(in):: kmax
            type(pblContrGradDataType), intent(out):: cbl

            allocate(cbl%entr(kmax))
            allocate(cbl%theta_up(kmax))
            allocate(cbl%thetav_up(kmax))
            allocate(cbl%qv_up(kmax))
            allocate(cbl%ua_up(kmax))
            allocate(cbl%va_up(kmax))
            allocate(cbl%wu(kmax))
        end subroutine cbl_allocate

        subroutine cbl_deallocate(cbl)
            type(pblContrGradDataType), intent(inout):: cbl

            deallocate(cbl%entr)
            deallocate(cbl%theta_up)
            deallocate(cbl%thetav_up)
            deallocate(cbl%qv_up)
            deallocate(cbl%ua_up)
            deallocate(cbl%va_up)
            deallocate(cbl%wu)
        end subroutine cbl_deallocate
        subroutine cbl_get_entrainment(cbl, bl, fluid,  grid)
            use scm_state_data, only : stateBLDataType
            use pbl_turb_data, only : turbBLDataType
            use phys_fluid, only: fluidParamsDataType
            use pbl_grid, only : pblgridDataType
            implicit none
            type(pblContrGradDataType) :: cbl
            type(stateBLDataType), intent(inout):: bl
            type(fluidParamsDataType), intent(in) :: fluid
            type(pblgridDataType), intent(in) :: grid

            integer k
            real dz

            cbl%entr(:) = 0.0
            do k = grid%kmax, 2, -1
                dz =(grid%z_cell(k-1) - grid%z_cell(k)) + 1.0
                if (k  < bl%kpbl+2) then
                    cbl%entr(k) = 0.0

                else
                    cbl%entr(k) = c_entr * ( 1.0/(grid%z_cell(k) + dz)  &
                            + 1.0/(bl%hpbla_diag - grid%z_cell(k) + dz) &
                            + 1.0/ lmin_entr)

                end if
                write(*,*) 'entr, k = ' , cbl%entr(k), k, grid%z_cell(k)
            end do
            cbl%entr(1) = 0
        end subroutine cbl_get_entrainment

        subroutine cbl_get_up(cbl, bl, fluid, grid)
            use scm_state_data, only : stateBLDataType
            use pbl_turb_data, only : turbBLDataType
            use phys_fluid, only: fluidParamsDataType
            use pbl_grid, only : pblgridDataType
            implicit none
            type(pblContrGradDataType) :: cbl
            type(stateBLDataType), intent(inout):: bl
            type(fluidParamsDataType), intent(in) :: fluid
            type(pblgridDataType), intent(in) :: grid
            real phys_beta, wstr, dz, dzentrp, dzentrm, temp
            integer k, kmax

            cbl%theta_up(:) = 0.0
            cbl%thetav_up(:) = 0
            cbl%qv_up = 0.0
            cbl%ua_up = 0.0
            cbl%va_up = 0.0

            kmax = grid%kmax
            phys_beta = 0.5 * fluid%g /(bl%theta_v(kmax) + bl%theta_v(kmax-1))
            write(*,*) 'bl%hs, ', (bl%surf%hs), phys_beta
            if (bl%surf%hs > 0.0) then
                wstr = (phys_beta * (bl%surf%hs/bl%rho(kmax)) &
                        * bl%hpbla_diag)**(0.33333)
                wstr = max(wstrmin, wstr)

                cbl%thetav_up0 = (bl%surf%hs/(bl%rho(kmax)*fluid%cp)) / wstr
                cbl%theta_up0 =  (bl%surf%hs/(bl%rho(kmax)*fluid%cp)) / wstr
                cbl%qv_up0 =  (bl%surf%es/bl%rho(kmax)) / wstr
            else
                cbl%thetav_up0 = 0.0
                cbl%theta_up0 = 0.0
                cbl%qv_up0 = 0.0
            end if
            cbl%ua_up0 =  0.0
            cbl%va_up0  =  0.0

            cbl%thetav_up(kmax) = cbl%thetav_up0 + bl%theta_v(kmax)
            cbl%theta_up(kmax) = cbl%theta_up0  + bl%theta(kmax)
            cbl%qv_up(kmax) = cbl%qv_up0  + bl%qv(kmax)
            cbl%ua_up(kmax) = cbl%ua_up0  + bl%u(kmax)
            cbl%va_up(kmax) = cbl%va_up0  + bl%v(kmax)

            do k = kmax-1 ,1,-1
                if (k <= bl%kpbl ) then
                    cbl%theta_up(k) = bl%theta(k)
                    cbl%thetav_up(k) = bl%theta_v(k)
                    cbl%qv_up(k) = bl%qv(k)
                    cbl%ua_up(k) = bl%u(k)
                    cbl%ua_up(k) = bl%u(k)
                    cbl%va_up(k) = bl%v(k)
                else
                    dz = (grid%z_cell(k) - grid%z_cell(k+1))
                    cbl%thetav_up(k) = cbl%thetav_up(k+1) - &
                            dz * cbl%entr(k+1) &
                                    * (cbl%thetav_up(k+1) - bl%theta_v(k+1))
                    cbl%theta_up(k) = cbl%theta_up(k+1) - &
                            dz * cbl%entr(k+1) &
                                    * (cbl%theta_up(k+1) - bl%theta(k+1))
                    write(*,*) 'theta_up, k = ' , cbl%theta_up(k) - bl%theta(k), k
                    cbl%qv_up(k) = cbl%qv_up(k+1) - &
                            dz * cbl%entr(k+1) &
                                    * (cbl%qv_up(k+1) - bl%qv(k+1))
                    cbl%ua_up(k) = cbl%ua_up(k+1) + &
                            dz * cbl%entr(k+1) &
                                    * (cbl%ua_up(k+1) - bl%u(k+1))
                    cbl%va_up(k) = cbl%va_up(k+1) - &
                            dz * cbl%entr(k+1) &
                                    * (cbl%va_up(k+1) - bl%v(k+1))
                end if
            end do

        end subroutine cbl_get_up

    subroutine cbl_get_wu(cbl, bl, fluid, grid)
        use scm_state_data, only : stateBLDataType
        use pbl_turb_data, only : turbBLDataType
        use phys_fluid, only: fluidParamsDataType
        use pbl_grid, only : pblgridDataType
        implicit none
        type(pblContrGradDataType) :: cbl
        type(stateBLDataType), intent(inout):: bl
        type(fluidParamsDataType), intent(in) :: fluid
        type(pblgridDataType), intent(in) :: grid
        real wu, phys_beta, umod, ustr, rhs, a, b, dz
        integer k, kmax
        kmax = grid%kmax
        cbl%wu(:) = 0.0


        umod = sqrt(bl%u(kmax)*bl%u(kmax) + bl%v(kmax)*bl%v(kmax))
        ustr = sqrt(bl%surf%cm2u /bl%rho(kmax) * umod)
        phys_beta = 0.5 * fluid%g /(bl%theta_v(kmax) + bl%theta_v(kmax-1))
        cbl%wu(kmax) = 2.46 * ((ustr)**3   &
                    - 0.4 * phys_beta * bl%surf%hs  &
                        * (grid%z_cell(k))/bl%rho(kmax))**(1.0/3.0) &
                            * sqrt(1.0 - grid%z_cell(k) / bl%hpbla_diag)
        do k = kmax - 1, 1, -1
            dz = (grid%z_cell(K) - grid%z_cell(K+1))
            a = 0.0
            if (abs(cbl%entr(k)) >= 1.0e-8 .and. abs(cbl%wu(k+1)) > 1.0e-6 ) then
                a = b1 * cbl%entr(k) * cbl%wu(k+1) * dz
                b =  b2 * 9.81 * (dz/cbl%wu(k+1))  &
                        * (cbl%thetav_up(k) - bl%theta_v(k))/bl%theta_v(k);

                rhs = cbl%wu(k+1) -  a  + b;
                if (rhs > 0) then
                    cbl%wu(k) = rhs;
                else
                    cbl%wu(k) = 0.0;
                endif
            else
                cbl%wu(k) = 0.0;
            endif
        enddo
        end subroutine cbl_get_wu

        subroutine cbl_apply_cntrg(cbl, bl, fluid, grid, dt)
            use scm_state_data, only : stateBLDataType
            use pbl_turb_data, only : turbBLDataType
            use phys_fluid, only: fluidParamsDataType
            use pbl_grid, only : pblgridDataType
            implicit none
            type(pblContrGradDataType), intent(inout) :: cbl
            type(stateBLDataType), intent(inout):: bl
            type(fluidParamsDataType), intent(in) :: fluid
            type(pblgridDataType), intent(in) :: grid
            real, intent(in) :: dt

            real, dimension(grid%kmax) :: rhs_up, rhs

            real a0w0
            integer kmax, kpbl, k
            kmax = grid%kmax
            kpbl = bl%kpbl
            a0w0 =  0.15 * cbl%wu(kpbl)


            call cbl_get_entrainment(cbl, bl, fluid,  grid)
            call cbl_get_up(cbl, bl, fluid, grid)
            call cbl_get_wu(cbl, bl, fluid, grid)
            !apply upwind
            rhs_up(:) = 0
            do k = kmax-1, 1, -1
                if (abs(cbl%wu(k)) > 1.0e-7) then
                    rhs_up(k) =  0.1 * (cbl%wu(k)) *  &
                            (cbl%theta_up(k+1) - cbl%theta_up(k) - bl%theta(k+1) + bl%theta(k))  &
                            / (grid%z_cell(k) - grid%z_cell(k+1))
                else
                    rhs_up(k)=0
                end if
            end do
            bl%theta = bl%theta + dt * rhs_up

            rhs_up = 0
            do k = kmax-1, 1, -1
                if (abs(cbl%wu(k)) > 1.0e-7) then
                    rhs_up(k) =  0.1 * (cbl%wu(k)) *  &
                            (cbl%qv_up(+1) - cbl%qv_up(k)  &
                                    - bl%qv(k+1) + bl%qv(k)) / (grid%z_cell(k) -grid%z_cell(k+1))
                else
                    rhs_up(k)=0
                end if
            end do
            bl%qv(:) = bl%qv(:) + dt * rhs_up(:)

            rhs_up(:) = 0
            do k = kmax-1, 1, -1
                if (abs(cbl%wu(k)) > 1.0e-8) then
                    rhs_up(k) =  0.1 * (cbl%wu(k)) *  &
                            (cbl%ua_up(k+1) - cbl%ua_up(k) - bl%u(k+1) + bl%u(k))  &
                            / (grid%z_cell(k) - grid%z_cell(k+1))
                else
                    rhs_up(k)=0
                end if
            end do
            bl%u(:) = bl%u(:) + dt * rhs_up(:)
            rhs_up(:) = 0
            do k = kmax-1, 1, -1
                if (abs(cbl%wu(k)) > 1.0e-8) then
                    rhs_up(k) =  0.1 * (cbl%wu(k)) *  &
                            (cbl%va_up(k+1) - cbl%va_up(k) - bl%v(k+1) + bl%v(k))  &
                            / (grid%z_cell(k) - grid%z_cell(k+1))
                else
                    rhs_up(k)=0
                end if
            end do
            bl%v(:) = bl%v(:) + dt * rhs_up(:)

            end subroutine cbl_apply_cntrg

    subroutine new_cntrg(cbl, bl, fluid, grid, dt)
        use scm_state_data, only : stateBLDataType
        use pbl_turb_data, only : turbBLDataType
        use phys_fluid, only: fluidParamsDataType
        use pbl_grid, only : pblgridDataType
        implicit none
        type(pblContrGradDataType), intent(inout) :: cbl
        type(stateBLDataType), intent(inout):: bl
        type(fluidParamsDataType), intent(in) :: fluid
        type(pblgridDataType), intent(in) :: grid
        real, intent(in) :: dt

        ! local variables
        real, save, allocatable ::  rhs_up(:), rhs(:)
        real, save, allocatable :: bouf(:)
        real wu, phys_beta, umod, rhs_ws, a, b
        real   dz, dzentrp, dzentrm, temp,dth
        real sigmaw, sigmath, sigmaqc
        real wstr, ustr, thstar, thvstar, qstar
        real z1, hpb, hpbmf
        real entrint, entexp, entexpu
        real wp, entw, wn2


        real a0w0
        integer kmax, kpbl, k
        kmax = grid%kmax
        kmax = grid%kmax
        if (.not.(allocated(rhs_up))) then
            allocate(rhs_up(grid%kmax), source=0.0)
        end if
        if (.not.(allocated(rhs))) then
            allocate(rhs(grid%kmax), source=0.0)
        end if
        if (.not.(allocated(bouf))) then
            allocate(bouf(grid%kmax), source=0.0)
        end if
        kpbl = bl%kpbl
        cbl%theta_up(:) = 0.0
        cbl%thetav_up(:) = 0
        cbl%qv_up(:) = 0.0
        cbl%ua_up(:) = 0.0
        cbl%va_up(:) = 0.0
        cbl%entr(:) = 0.0
        cbl%wu(:) = 0.0
        rhs_up(:) = 0.0
        rhs(:) = 0.0
        ! Check if the flux is convective
        if (bl%surf%hs > 0.0) then
            ! get surface properties
            z1  = grid%z_cell(kmax)
            hpb = bl%hpbla_diag
            phys_beta = 0.5 * fluid%g /(bl%theta_v(kmax) + bl%theta_v(kmax-1))
            wstr = (phys_beta * (bl%surf%hs/bl%rho(kmax)/fluid%cp) &
                    * bl%hpbla_diag)**(0.33333)
            wstr = max(wstrmin, wstr)

            ! Initial condition for vertical velocity
            umod = sqrt(bl%u(kmax)*bl%u(kmax) + bl%v(kmax)*bl%v(kmax))
            ustr = sqrt(bl%surf%cm2u /bl%rho(kmax) * umod)
            !cbl%wu(kmax) = 1.0 * ((ustr)**3.0   & !2.46
            !        + 0.4 * wstr**3.0)**(1.0/3.0) &
            !        * sqrt(1.0 - grid%z_cell(kmax) / bl%hpbla_diag)

            sigmaw = 1.3 * (ustr**3.0 + 0.6 * wstr**3.0 * z1/hpb) * &
                    sqrt(1.0 - z1 / hpb)

            thstar = (bl%surf%hs/(bl%rho(kmax))) / wstr
            thvstar =  (bl%surf%hs/(bl%rho(kmax))) / wstr
            qstar = (bl%surf%es/bl%rho(kmax)) / wstr
            cbl%wu(kmax) = 0.5 * sigmaw


            ! initial perturbation
            cbl%thetav_up0 = (bl%surf%hs/(bl%rho(kmax))) / sigmaw
            cbl%theta_up0 =  (bl%surf%hs/(bl%rho(kmax))) / sigmaw
            cbl%qv_up0 =  (bl%surf%es/bl%rho(kmax)) / sigmaw
            cbl%ua_up0 =  0.5 * (bl%u(kmax) + bl%u(kmax-1))
            cbl%va_up0  =  0.5 *  (bl%v(kmax) + bl%v(kmax-1))
            cbl%thetav_up(kmax) =  alph0 *cbl%thetav_up0 + bl%theta_v(kmax)
            cbl%theta_up(kmax) = alph0 * cbl%theta_up0  + bl%theta(kmax)
            cbl%qv_up(kmax) = alph0 * cbl%qv_up0  + bl%qv(kmax)
            cbl%ua_up(kmax) = cbl%ua_up0  + bl%u(kmax)
            cbl%va_up(kmax) = cbl%va_up0  + bl%v(kmax)



            !Initial entrainment
            !cbl%entr(kmax) = 1.0 / (cbl%wu(kmax) * tau_dry_scale)  &
                            !+ c_entr / grid%z_cell(kmax)
            dz = grid%z_cell(kmax) - grid%z_edge(kmax)
            cbl%entr(kmax) = c_entr * ( 1.0/(grid%z_cell(kmax) + dz)  &
                    + 1.0/(bl%hpbla_diag - grid%z_cell(kmax) + dz))

            bl%kpbld_mf = bl%kpbl
            ! get non-entraiened parcel height
            !do k = kmax - 1, 2, -1
            !    dth = cbl%thetav_up(kmax) - bl%theta_v(k)
            !    if (bl%kpbld_mf == 0 .and. dth < 0.0 ) then
            !        bl%kpbld_mf = k
            !    end if
            !end do
            hpbmf = 0.5 * (grid%z_cell(bl%kpbld_mf) + &
                    bl%hpbla_diag)
            write(*,*) 'hbp, hmf', grid%z_cell(bl%kpbld_mf), bl%hpbla_diag, bl%surf%hs
            ! get proper entrainment
            do k = grid%kmax-1,1,-1
                dz = (grid%z_cell(k) - grid%z_cell(k+1))
                if (grid%z_cell(k) <= hpbmf) then
                cbl%entr(k) = c_entr * ( 1.0/(grid%z_cell(k) + dz)  &
                        + 1.0/(hpbmf - grid%z_cell(k) + dz))
                else
                    cbl%entr(k) = 0.0
                end if
            end do

            !get updrafts
            do k = grid%kmax-1,bl%kpbld_mf,-1
                dz = (grid%z_cell(k) - grid%z_cell(k+1))
                entexp = exp( - cbl%entr(k) * dz)
                entexpu = exp( - cbl%entr(k) * dz / 3.0 )
                cbl%thetav_up(k) = bl%theta_v(k) * (1.0 - entexp) + &
                        cbl%thetav_up(k+1) * entexp
                cbl%theta_up(k) = bl%theta(k) * (1.0 - entexp) + &
                        cbl%theta_up(k+1) * entexp
                cbl%qv_up(k) = bl%qv(k) * (1.0 - entexp) + &
                        cbl%qv_up(k+1) * entexp
                cbl%ua_up(k) = bl%u(k) * (1.0 - entexpu) + &
                        cbl%ua_up(k+1) * entexpu
                cbl%va_up(k) = bl%u(k) * (1.0 - entexpu) + &
                        cbl%va_up(k+1) * entexpu
                bouf(k) = fluid%g *  &
                     (0.5 * (cbl%thetav_up(k) + cbl%thetav_up(k+1)) / &
                        bl%theta_v(k) - 1.0)
            enddo

            !get vertical speed
            do k = grid%kmax-1,bl%kpbld_mf,-1
                dz = (grid%z_cell(k) - grid%z_cell(k+1))
                wp = b1 * cbl%entr(k)
                if (wp > 0.0) then
                    entw = exp( -2.0 * wp * dz )
                    wn2 = cbl%wu(k+1) * cbl%wu(k+1) * entw + &
                            b2* bouf(k)*(1.0-entw)/entw
                else
                    wn2 = cbl%wu(k+1) * cbl%wu(k+1) + &
                            2.0 * b2 * bouf(k) * dz
                end if

                if ( wn2 > 0.0) then
                    cbl%wu(k) = sqrt(wn2)
                else
                    cbl%wu(k) = 0.0
                end if
            end do


            ! apply upwind
            rhs_up(:) = 0.0
            do k = kmax-1, 1, -1
                if (cbl%wu(k) > 0.0) then
                    rhs_up(k) =  a0 * (cbl%wu(k)) *  &
                            (cbl%theta_up(k+1) - cbl%theta_up(k) - bl%theta(k+1) + bl%theta(k))  &
                            / (grid%z_cell(k) - grid%z_cell(k+1))
                else
                    rhs_up(k)=0.0
                end if
            end do
            bl%theta = bl%theta + dt * rhs_up

            rhs_up(:) = 0.0
            do k = kmax-1, 1, -1
                if (cbl%wu(k) > 0.0) then
                    rhs_up(k) =  a0 * (cbl%wu(k)) *  &
                            (cbl%qv_up(k+1) - cbl%qv_up(k)  &
                                    - bl%qv(k+1) + bl%qv(k)) &
                            / (grid%z_cell(k) -grid%z_cell(k+1))
                else
                    rhs_up(k)=0.0
                end if
            end do
            bl%qv(:) = bl%qv(:) + dt * rhs_up(:)

            rhs_up(:) = 0.0
            do k = kmax-1, 1, -1
                if (cbl%wu(k) > 0.0) then
                    rhs_up(k) =  a0 * (cbl%wu(k)) *  &
                            (cbl%ua_up(k+1) - cbl%ua_up(k) - bl%u(k+1) + bl%u(k))  &
                            / (grid%z_cell(k) - grid%z_cell(k+1))
                else
                    rhs_up(k)=0.0
                end if
            end do
            bl%u(:) = bl%u(:) + dt * rhs_up(:)
            rhs_up(:) = 0.0
            do k = kmax-1, 1, -1
                if (cbl%wu(k) > 0.0) then
                    rhs_up(k) =  a0 * (cbl%wu(k)) *  &
                            (cbl%va_up(k+1) - cbl%va_up(k) - bl%v(k+1) + bl%v(k))  &
                            / (grid%z_cell(k) - grid%z_cell(k+1))
                else
                    rhs_up(k)=0.0
                end if
            end do
            bl%v(:) = bl%v(:) + dt * rhs_up(:)
        end if
    end subroutine new_cntrg
end module pbl_dry_contrgradient