! Created by Andrey Debolskiy on 11.02.2025.

!> Collection of utilities for pbl_turb_data
module pbl_turb_common
use pbl_turb_data
implicit none
    public get_s2, get_n2, get_rig
    public get_coeffs_general

contains
    subroutine get_n2(turb, bl, fluid, grid)
        use pbl_turb_data, only: turbBLDataType
        use scm_state_data, only: stateBLDataType
        use pbl_grid, only: pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        implicit none
        type(fluidParamsDataType), intent(in):: fluid
        type(stateBLDataType), intent(in):: bl
        type(pblgridDataType), intent(in)::grid
        type(turbBLDataType), intent(inout):: turb

        real dtvir, buoyp
        integer k
        do  k = 1,grid%kmax-1
            dtvir = bl%theta_v(k) - bl%theta_v(k+1)
            buoyp = fluid%g / ( (bl%theta_v(k+1) + bl%theta_v(k))/2.0e0 )
            turb%n2 = buoyp*(grid%z_cell(k) - grid%z_cell(k+1))*dtvir
        end do
    end subroutine get_n2

    subroutine get_s2(turb, bl, fluid, grid, du2min)
        use pbl_turb_data, only: turbBLDataType
        use scm_state_data, only: stateBLDataType
        use pbl_grid, only: pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        implicit none
        type(fluidParamsDataType), intent(in):: fluid
        type(stateBLDataType), intent(in):: bl
        type(pblgridDataType), intent(in)::grid
        type(turbBLDataType), intent(inout):: turb
        real, intent(in) :: du2min

        real du2
        integer k
        do  k = 1,grid%kmax-1
            du2 = amax1(du2min, &
                    &    (bl%u(k)-bl%u(k+1))**2 + (bl%v(k)-bl%v(k+1))**2)
            turb%s2 = du2 / ((grid%z_cell(k) - grid%z_cell(k+1)) &
                        * (grid%z_cell(k) - grid%z_cell(k+1)))
        enddo
    end subroutine get_s2

    subroutine get_rig(turb, bl, fluid, grid, du2min, ricr)
    use pbl_turb_data, only: turbBLDataType
    use scm_state_data, only: stateBLDataType
    use pbl_grid, only: pblgridDataType
    use phys_fluid, only: fluidParamsDataType
    implicit none
    type(fluidParamsDataType), intent(in):: fluid
    type(stateBLDataType), intent(in):: bl
    type(pblgridDataType), intent(in)::grid
    real, intent(in):: du2min, ricr
    type(turbBLDataType), intent(inout):: turb

    real dtvir, buoyp, du2, rinum
    integer k
    do  k = 1,grid%kmax-1
        du2 = amax1(du2min, &
                &    (bl%u(k)-bl%u(k+1))**2 + (bl%v(k)-bl%v(k+1))**2)
        dtvir = bl%theta_v(k) - bl%theta_v(k+1)
        buoyp = fluid%g / ( (bl%theta_v(k+1) + bl%theta_v(k))/2.0e0 )
        rinum = buoyp*(grid%z_cell(k) - grid%z_cell(k+1))*dtvir/du2
        rinum = amin1(ricr,rinum)
        turb%rig(k) = rinum
    end do
    end subroutine get_rig

    subroutine get_coeffs_general(turb, bl, fluid, hbl_option,  grid)
        use pbl_turb_data, only: turbBLDataType
        use scm_state_data, only: stateBLDataType
        use pbl_grid, only: pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        use pbl_fo_turb, only: get_fo_coeffs, get_turb_length, du2min, ricr
        implicit none
        type(fluidParamsDataType), intent(in):: fluid
        type(stateBLDataType), intent(inout):: bl
        type(pblgridDataType), intent(in)::grid
        type(turbBLDataType), intent(inout):: turb
        integer, intent(in):: hbl_option


        call  get_n2(turb, bl, fluid, grid)

        call get_s2(turb, bl, fluid, grid, du2min)

        call get_rig(turb, bl, fluid, grid, du2min, ricr)
        call  get_turb_length(turb, bl, fluid, grid, hbl_option)

        if (turb%cl_type == 1) then ! FOM closures

            call get_fo_coeffs(turb, bl, fluid,  grid)
        end if
        if (turb%cl_type == 2) then ! KL closures
            !call get_kl_coeffs(turb, bl, fluid,  grid)
        end if

    end subroutine get_coeffs_general
end module pbl_turb_common