MODULE vertical_remap_mod
   USE prec, ONLY: rstd

   IMPLICIT NONE
   PRIVATE

   PUBLIC vertical_remap, vertical_remap_extdata, compute_vertical_remap_extdata

CONTAINS

   SUBROUTINE vertical_remap(pressure_level, field_in, f_ps, field_out)
      USE icosa
      USE omp_para
      USE pression_mod, ONLY: pression
      REAL(rstd), INTENT(IN) :: pressure_level(:)
      TYPE(t_field), POINTER :: field_in(:)
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: field_out(:)

      TYPE(t_field), POINTER, SAVE :: f_p(:)
      REAL(rstd), POINTER :: in(:, :)
      REAL(rstd), POINTER :: out(:, :)
      REAL(rstd), POINTER :: p(:, :)

      INTEGER :: ind, iq

      CALL allocate_field(f_p, field_t, type_real, llm + 1)
      CALL pression(f_ps, f_p)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         p = f_p(ind)
         DO iq = 1, field_out(ind)%dim4
            in => field_in(ind)%rval4d(:, :, iq)
            out => field_out(ind)%rval4d(:, :, iq)
            CALL compute_vertical_remap(pressure_level, in, p, out)
         END DO
      END DO
      CALL deallocate_field(f_p)

   END SUBROUTINE vertical_remap

   SUBROUTINE compute_vertical_remap(pressure_level, in, p, out)
      USE icosa
      USE omp_para
      REAL(rstd), INTENT(IN)  :: pressure_level(:)
      REAL(rstd), INTENT(IN)  :: in(:, :)
      REAL(rstd), INTENT(IN)  :: p(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT) :: out(iim*jjm, llm)
      REAL(rstd) :: coeff, pmid
      INTEGER :: i, j, ij, l, n, nb_level
      INTEGER :: a
      INTEGER :: b
      LOGICAL :: positive

      nb_level = size(pressure_level)
      IF (pressure_level(1) >= pressure_level(nb_level)) THEN
         positive = .FALSE.
      ELSE
         positive = .TRUE.
      END IF

      !$OMP BARRIER
      IF (is_omp_level_master) THEN

         DO l = 1, llm
            DO j = jj_begin, jj_end
               DO i = ii_begin, ii_end
                  ij = (j - 1)*iim + i
                  pmid = 0.5*(p(ij, l) + p(ij, l + 1))
                  IF (positive) THEN
                     a = 0
                     DO n = 1, nb_level - 1
                        IF ((pmid >= pressure_level(n) .AND. pmid < pressure_level(n + 1))) THEN
                           a = n; b = n + 1; EXIT
                        END IF
                     END DO
                     IF (a == 0) THEN
                        IF (pmid <= pressure_level(1)) THEN
                           a = 1; b = 2
                        ELSE
                           a = nb_level - 1; b = nb_level
                        END IF
                     END IF
                  ELSE
                     a = 0
                     DO n = 1, nb_level - 1
                        IF ((pmid <= pressure_level(n) .AND. pmid > pressure_level(n + 1))) THEN
                           a = n; b = n + 1; EXIT
                        END IF
                     END DO

                     IF (a == 0) THEN
                        IF (pmid >= pressure_level(1)) THEN
                           a = 1; b = 2
                        ELSE
                           a = nb_level - 1; b = nb_level
                        END IF
                     END IF
                  END IF

                  coeff = (pmid - pressure_level(a))/(pressure_level(a) - pressure_level(b))
                  out(ij, l) = in(ij, a) + coeff*(in(ij, a) - in(ij, b))
               END DO
            END DO
         END DO

      END IF
      !$OMP BARRIER

   END SUBROUTINE compute_vertical_remap

   SUBROUTINE vertical_remap_extdata(field_in, f_target_pressure, field_out)
      USE icosa
      USE omp_para
      USE disvert_mod, ONLY: presnivs

      IMPLICIT NONE
      TYPE(t_field), POINTER :: field_in(:)
      TYPE(t_field), POINTER :: field_out(:)
      TYPE(t_field), POINTER :: f_target_pressure(:)

      REAL(rstd), POINTER :: target_pressure(:, :)
      REAL(rstd), POINTER :: in(:, :)
      REAL(rstd), POINTER :: out(:, :)
      INTEGER :: ind

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         in = field_in(ind)
         out = field_out(ind)
         target_pressure = f_target_pressure(ind)
         CALL compute_vertical_remap_extdata(in, target_pressure, out)
      END DO

   END SUBROUTINE vertical_remap_extdata

   SUBROUTINE compute_vertical_remap_extdata(in, target_pressure, out2d_press)
      USE icosa
      USE omp_para
      USE disvert_mod, ONLY: presnivs
      IMPLICIT NONE
      REAL(rstd), INTENT(IN)  :: in(:, :)
      REAL(rstd), INTENT(IN)  :: target_pressure(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: out2d_press(iim*jjm, llm)
      REAL(rstd) :: coeff, target_pval, testp1, testp2
      INTEGER :: i, j, ij, l, n, nb_level
      INTEGER :: a
      INTEGER :: b
      LOGICAL :: positive

      nb_level = size(presnivs)
      !$OMP BARRIER
      IF (is_omp_level_master) THEN
      DO l = 1, llm
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               target_pval = target_pressure(ij, l)
               a = 0
               DO n = 1, nb_level - 1
                  IF ((target_pval <= presnivs(n) .AND. target_pval >= presnivs(n + 1))) THEN
                     testp1 = presnivs(n); testp2 = presnivs(n + 1)
                     a = n; b = n + 1; EXIT
                  END IF
               END DO
               IF (a == 0) THEN
                  IF (target_pval >= presnivs(1)) THEN
                     a = 1; b = 2
                  ELSE
                     a = nb_level - 1; b = nb_level
                  END IF
               END IF

               coeff = (target_pval - presnivs(a))/(presnivs(a) - presnivs(b))
               out2d_press(ij, l) = in(ij, a) + coeff*(in(ij, a) - in(ij, b))
            END DO
         END DO
      END DO

      END IF
      !$OMP BARRIER

   END SUBROUTINE compute_vertical_remap_extdata

END MODULE vertical_remap_mod
