! Created by Andrey Debolskiy on 17.12.2023.

module state_utils

    public geo
    public get_geopotential_moist
    public theta2ta
    public ta2theta
    public get_exner
    public get_sigma_from_z
    public get_z_from_sigma
    public get_z_from_sigma_theta
    public get_density_theta
    public get_density
    public dqsdt_sc, qmax_sc, qmax
    public DQSDT2, DQSDT, CLDT
    public get_theta_v
    contains
        subroutine get_geopotential_moist(F,t, qv, sig, fs, rd, eps, kmax)
            implicit none
            integer, intent(in):: kmax
            real, intent(in):: t(kmax)
            real, intent(in):: qv(kmax)
            real, intent(in):: sig(kmax)
            real, intent(in):: rd
            real, intent(in):: eps
            real, intent(in):: fs
            real, intent(out)::f(kmax)

            integer ka, k
            real thalf, qhalf

            F(kmax) = fs - log(sig(kmax)) * rd * t(kmax) * (1.0 + eps * qv(kmax)) / (1.0 + qv(kmax))

            do k = kmax-1, 1, -1
                ka = kmax+1
                thalf = 0.5 * (t(ka) + t(k))
                qhalf = 0.5 * (qv(ka) + qv(k))
                F(k) = F(ka) - (log(sig(ka)) -log(sig(k))) &
                    * rd * thalf * (1.0 + eps * qhalf) / (1.0 + qhalf)
            end do
        end subroutine get_geopotential_moist

        SUBROUTINE GEO(T,FS,F,SIG, R, LL)
        implicit none
        integer, intent(in):: LL
        real, intent(in):: T(LL)
        real, intent(in):: SIG(LL)
        real, intent(in):: FS
        real, intent(out):: F(LL)
        real, intent(in):: R

        integer k, ka

        F(LL)=FS-ALOG(SIG(LL))*R*T(LL)

        do  K=LL-1,1,-1

            KA=K+1
            F(K)=F(KA)+0.5E0*R*(ALOG(SIG(KA))-ALOG(SIG(K)))* &
             (T(KA)+T(K))

        end do

        END SUBROUTINE GEO




        subroutine theta2ta(ta, theta, p0, sig, appa, kl)
            implicit none
            real, dimension(kl), intent(in):: theta
            real, dimension(kl), intent(in)::sig
            real, intent(in):: p0, appa
            integer, intent(in):: kl
            real, dimension(kl), intent(out):: ta

            integer k

            do k=1,kl
                ta(k) = theta(k) * (p0 * sig(k))**appa
            end do
        end subroutine theta2ta

        subroutine ta2theta(theta, ta, p0, sig, appa, kl)
            implicit none
            real, dimension(kl), intent(in):: ta
            real, dimension(kl), intent(in)::sig
            real, intent(in):: p0, appa
            integer, intent(in):: kl
            real, dimension(kl), intent(out):: theta

            integer k

            do k=1,kl
                theta(k) = ta(k) * (p0 * sig(k))**(-appa)
            end do
        end subroutine ta2theta

        real function qmax_sc(t, p, aa)
            implicit none

            real, intent(in) :: t, p, aa
            real :: pb2, pb

            pb2 = t - 273.2e0
            if ((aa.gt.1.0e0.and.pb2.ge.0.0e0).or.(aa.lt.1.0e0.and.pb2.ge.-12.0e0)) then
                pb=17.55e0/(t-31.1e0)
            else
                pb=21.85e0/(t-7.65e0)
            end if
            qmax_sc=3.80042e-3*exp(pb*pb2)/p
        end function qmax_sc
    real function qmax(T,P, AA )
        IMPLICIT NONE
        real, intent(in):: t, p, aa
        !local
        real pb2, pb
        PB2=T-273.2E0
        IF((AA.GT.1.0E0.AND.PB2.GE.0.0E0).OR. &
                 (AA.LT.1.0E0.AND.PB2.GE.-12.0E0))THEN
            PB=17.55E0/(T-31.1E0 )
        ELSE
            PB=21.85E0/(T-7.65E0)
        ENDIF
        QMAX=3.80042E-3*EXP(PB*PB2)/P
        RETURN
    end function qmax


    real function dqsdt_sc(a, b, h)
            real, intent(in) ::  a, b, h

            dqsdt_sc=4248.855e0*h*qmax_sc(a,b, 0.0e0)/(a-31.10e0)**2
        end function dqsdt_sc

    REAL FUNCTION CLDT(Q,T,P,C)
        !
        IMPLICIT NONE
        REAL, INTENT(IN) ::  Q,T,P,C
        !      REAL DQSDT
        !REAL :: QMAX

        CLDT=(Q-C*QMAX(T,P, 0.0E0))/(1.0E0+2488.851E0*DQSDT(T,P,C))
        RETURN
    END FUNCTION CLDT

    REAL FUNCTION DQSDT(A,B,H)
        !
        IMPLICIT NONE
        REAL, INTENT(IN) ::  A, B, H
        !REAL :: QMAX

        DQSDT=4248.855E0*H*QMAX(A,B, 0.0E0)/(A-31.10E0)**2
        RETURN
    END FUNCTION DQSDT

    REAL FUNCTION DQSDT2(A,B,H)

        IMPLICIT NONE

        REAL, INTENT(IN) ::  A, B, H
        !      REAL  DQSDT
        !REAL :: QMAX

        DQSDT2=4248.855E0*H*(DQSDT(A,B,    H )/(A-31.10E0)**2 - &
                2.0E0*QMAX (A,B, 0.0E0)/(A-31.10E0)**3)
        RETURN
    END FUNCTION DQSDT2

    real function get_density( p, t, qv, fluid)
        use phys_fluid, only: fluidParamsDataType
        implicit none
        type(fluidParamsDataType),intent(in):: fluid
        real, intent(in):: p, t, qv

        get_density = (p / (fluid%R_d * t)) * (1.0 + qv) / (1 + qv * fluid%eps_v)
    end function get_density

    real function get_density_theta( p, th, qv, fluid)
        use phys_fluid, only: fluidParamsDataType
        implicit none
        type(fluidParamsDataType),intent(in):: fluid
        real, intent(in):: p, th, qv

        real t
        t =  th * (p / fluid%p0)**(fluid%R_d / fluid%cp)
        get_density_theta = (p * 100.0 / (fluid%R_d * t)) * (1.0 + qv) / (1 + qv * fluid%eps_v)
    end function get_density_theta

    subroutine get_density_theta_vector( bl, fluid, grid)
        use phys_fluid, only: fluidParamsDataType
        use scm_state_data, only: stateBLDataType
        use pbl_grid, only: pblgridDataType
        implicit none
        type(fluidParamsDataType),intent(in):: fluid
        type(pblgridDataType),intent(in):: grid
        type(stateBLDataType),intent(inout)::bl

        real t, p, qv
        integer k

        do k=grid%kmax, 1, -1
            p = bl%p0 * grid%sigma(k)
            t = bl%theta(k) * (p / fluid%p0)**(fluid%R_d / fluid%cp)
            qv = bl%qv(k)
            bl%rho(k) = (p * 100.0 / (fluid%R_d * t)) * (1.0 + qv) / (1 + qv * fluid%R_v/fluid%R_d)
        end do

    end subroutine get_density_theta_vector

    subroutine get_theta_v( theta_v, theta, qv,  fluid, kmax)
        use phys_fluid, only: fluidParamsDataType
        implicit none
        integer,intent(in):: kmax
        type(fluidParamsDataType), intent(in):: fluid
        real,dimension(kmax), intent(in):: theta
        real,dimension(kmax), intent(in):: qv
        real,dimension(kmax), intent(out):: theta_v

        integer k

        do k = kmax,1,-1
            theta_v(k) = theta(k) * (1.0 + fluid%R_d / fluid%R_v * qv(k))
        end do
    end subroutine get_theta_v

    subroutine get_exner( exnr, p, fluid, kmax)
        use phys_fluid, only: fluidParamsDataType
        implicit none
        integer,intent(in):: kmax
        type(fluidParamsDataType), intent(in):: fluid
        real,dimension(kmax), intent(in):: p
        real,dimension(kmax), intent(out):: exnr

        integer k

        do k = kmax,1,-1
            exnr(k) = (p(k)/fluid%p0)**(fluid%R_d / fluid%cp)
        end do
    end subroutine get_exner

    subroutine get_exner_sig( exnr, ps, fluid, grid, kmax)
        use pbl_grid, only : pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        implicit none
        integer,intent(in):: kmax
        type(fluidParamsDataType), intent(in):: fluid
        type(pblgridDataType), intent(in)::grid
        real, intent(in):: ps
        real,dimension(kmax), intent(out):: exnr

        integer k
        do k = kmax,1,-1
            exnr(k) = (ps * grid%sigma(k)/fluid%p0)**(fluid%R_d / fluid%cp)
        end do
    end subroutine get_exner_sig

    subroutine get_sigma_from_z( grid, bl, ps, fluid)
        use pbl_grid, only : pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        use scm_state_data, only: stateBLDataType
        implicit none
        type(fluidParamsDataType), intent(in):: fluid
        type(stateBLDataType), intent(in):: bl
        real, intent(in):: ps
        type(pblgridDataType), intent(inout)::grid


        real dp, rho, p
        integer k, kmax
        kmax = grid%kmax
        grid%sigma_edge(kmax) = 1.0
        rho = get_density(ps, bl%temp(kmax), bl%qv(kmax), fluid)

        dp = - rho * fluid%g * grid%dzc(kmax)
        grid%sigma(kmax) = 1.0 + (dp/2) / ps
        grid%sigma_edge(kmax) = 1.0
        p = ps + dp

        do k = kmax-1,1,-1
            rho = get_density(p, bl%temp(k), bl%qv(k), fluid)
            dp = - rho * fluid%g * grid%dzc(k)
            grid%sigma_edge(k) = grid%sigma_edge(k+1) + dp / ps
            rho = get_density(p, 0.5*(bl%temp(k+1) + bl%temp(k)), &
                    0.5*(bl%qv(k+1) + bl%qv(k)), fluid)
            grid%sigma(k) = grid%sigma(k+1)  &
                    - rho * fluid%g * grid%dze(k) / ps
            p = grid%sigma_edge(k) * ps
        end do
        grid%dsigmac(2:kmax) = grid%sigma(2:kmax) - grid%sigma(1:kmax-1)
        grid%sigma_edge(0) = 0.0
        grid%dsigmae(1:kmax) = grid%sigma_edge(1:kmax) - grid%sigma_edge(1:kmax-1)
    end subroutine get_sigma_from_z

    subroutine get_sigma_from_z_theta( grid, bl, ps, fluid)
        use pbl_grid, only : pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        use scm_state_data, only: stateBLDataType
        implicit none
        type(fluidParamsDataType), intent(in):: fluid
        type(stateBLDataType), intent(in):: bl
        real, intent(in):: ps
        type(pblgridDataType), intent(inout)::grid


        real dp, rho, p, temp, appa
        integer k, kmax
        kmax = grid%kmax
        grid%sigma_edge(kmax) = 1.0
        !convert theta -> temp
        appa = fluid%R_d / fluid%cp
        temp = bl%theta(kmax) * (100000.0 / ps )**appa
        rho = get_density(ps, temp, bl%qv(kmax), fluid)

        dp = - rho * fluid%g * grid%dzc(kmax)
        grid%sigma(kmax) = 1.0 + (dp/2) / ps
        grid%sigma_edge(kmax) = 1.0
        p = ps + dp

        do k = kmax-1,1,-1
            temp  = bl%theta(k) * (100000.0 / p )**appa
            rho = get_density(p, bl%temp(k), bl%qv(k), fluid)
            dp = - rho * fluid%g * grid%dzc(k)
            grid%sigma_edge(k) = grid%sigma_edge(k+1) + dp / ps
            rho = get_density(p, 0.5*(bl%temp(k+1) + bl%temp(k)), &
                    0.5*(bl%qv(k+1) + bl%qv(k)), fluid)
            grid%sigma(k) = grid%sigma(k+1)  &
                    - rho * fluid%g * grid%dze(k) / ps
            p = grid%sigma_edge(k) * ps
        end do
        grid%dsigmac(2:kmax) = grid%sigma(2:kmax) - grid%sigma(1:kmax-1)
        grid%sigma_edge(0) = 0.0
        grid%dsigmae(1:kmax) = grid%sigma_edge(1:kmax) - grid%sigma_edge(1:kmax-1)
    end subroutine get_sigma_from_z_theta

    subroutine get_z_from_sigma( grid, bl, zs, fluid, kmax)
        use pbl_grid, only : pblgridDataType
        use phys_fluid, only: fluidParamsDataType
        use scm_state_data, only: stateBLDataType
        implicit none
        integer,intent(in):: kmax
        type(fluidParamsDataType), intent(in):: fluid
        type(stateBLDataType), intent(in):: bl
        real, intent(in):: zs
        type(pblgridDataType), intent(inout)::grid

        real fs
        real, dimension(kmax):: f
        integer k
        call get_geopotential_moist(f,bl%temp, bl%qv, grid%sigma, &
                fluid%g * zs, fluid%R_d, fluid%eps_v, kmax)
        grid%z_cell(1:kmax) = f(1:kmax) / fluid%g
        grid%z_edge(kmax) = zs
        do k = kmax - 1, 0, -1
            grid%z_edge(k) = grid%z_cell(kmax + 1) + 0.5 *( grid%z_cell(kmax + 1) - grid%z_edge(kmax + 1))
        end do
        grid%dzc(2:kmax) = grid%z_cell(2:kmax) - grid%z_cell(1:kmax-1)
    end subroutine get_z_from_sigma
end module state_utils