MODULE set_bounds_mod
   USE geometry, ONLY: swap_geometry, lon_e, lat_e
   USE dimensions, ONLY: swap_dimensions, u_pos
   USE math_const, ONLY: Pi
   USE domain_mod, ONLY: t_domain, t_cellset, domloc_glo_ind
   USE spherical_geom_mod, ONLY: xyz2lonlat

   IMPLICIT NONE
   PRIVATE
   SAVE

   PUBLIC :: set_bounds

CONTAINS

   SUBROUTINE set_bounds_primal(d, cells, all, own)
      TYPE(t_domain) :: d
      TYPE(t_cellset) :: cells
      LOGICAL :: all, own(:, :)         ! if all is .TRUE., include halo cells
      REAL :: lon, lat
      INTEGER :: i, j, k, n, halo_size

      halo_size = MERGE(1, 0, all)

      ! count primal cells
      n = 0
      DO j = d%jj_begin - halo_size, d%jj_end + halo_size
         DO i = d%ii_begin - halo_size, d%ii_end + halo_size
            IF (own(i, j) .OR. all) n = n + 1
         END DO
      END DO
      cells%ncell = n

      IF (.NOT. ALLOCATED(cells%ij)) THEN
         ! now set bounds
         ALLOCATE (cells%ij(n), cells%lon(n), cells%lat(n), cells%ind_glo(n))
         ALLOCATE (cells%bnds_lon(0:5, n), cells%bnds_lat(0:5, n))
      END IF

      n = 0
      DO j = d%jj_begin - halo_size, d%jj_end + halo_size
         DO i = d%ii_begin - halo_size, d%ii_end + halo_size
            IF (own(i, j) .OR. all) THEN
               n = n + 1
               CALL xyz2lonlat(d%xyz(:, i, j), lon, lat)
               cells%lon(n) = lon*180./Pi
               cells%lat(n) = lat*180./Pi
               DO k = 0, 5
                  CALL xyz2lonlat(d%vertex(:, k, i, j), lon, lat)
                  cells%bnds_lon(k, n) = lon*180./Pi
                  cells%bnds_lat(k, n) = lat*180./Pi
               END DO
               cells%ij(n) = d%iim*(j - 1) + i
               cells%ind_glo(n) = d%assign_cell_glo(i, j) - 1
            END IF
         END DO
      END DO

!    PRINT *, 'set_bounds_primal', all, halo_size, cells%ncell
   END SUBROUTINE set_bounds_primal

   SUBROUTINE set_bounds_dual(d, cells)
      USE metric, ONLY: vup, vdown
      TYPE(t_domain) :: d
      TYPE(t_cellset) :: cells
      REAL :: lonc, latc, lon(0:2), lat(0:2)
      INTEGER :: i, j, k, n

      ! count dual cells
      n = 0
      DO j = d%jj_begin + 1, d%jj_end
         DO i = d%ii_begin, d%ii_end - 1
            n = n + 2
         END DO
      END DO
      cells%ncell = n

      IF (.NOT. ALLOCATED(cells%ind_glo)) THEN
         ! now set bounds
         ! We are not able to assign a global index to vertices => ind_glo not allocated
         ALLOCATE (cells%ij(n), cells%lon(n), cells%lat(n))
         ALLOCATE (cells%bnds_lon(0:2, n), cells%bnds_lat(0:2, n))
      END IF

      n = 0
      DO j = d%jj_begin + 1, d%jj_end
         DO i = d%ii_begin, d%ii_end - 1
            n = n + 1
            CALL xyz2lonlat(d%vertex(:, vdown, i, j), lonc, latc)
            CALL xyz2lonlat(d%xyz(:, i, j), lon(0), lat(0))
            CALL xyz2lonlat(d%xyz(:, i, j - 1), lon(1), lat(1))
            CALL xyz2lonlat(d%xyz(:, i + 1, j - 1), lon(2), lat(2))
            cells%lon(n) = lonc*180./Pi
            cells%lat(n) = latc*180/Pi
            DO k = 0, 2
               cells%bnds_lat(k, n) = lat(k)*180./Pi
               cells%bnds_lon(k, n) = lon(k)*180./Pi
            END DO
            cells%ij(n) = d%z_down + d%iim*(j - 1) + i
         END DO
      END DO

      DO j = d%jj_begin, d%jj_end - 1
         DO i = d%ii_begin + 1, d%ii_end
            n = n + 1
            CALL xyz2lonlat(d%vertex(:, vup, i, j), lonc, latc)
            CALL xyz2lonlat(d%xyz(:, i, j), lon(0), lat(0))
            CALL xyz2lonlat(d%xyz(:, i, j + 1), lon(1), lat(1))
            CALL xyz2lonlat(d%xyz(:, i - 1, j + 1), lon(2), lat(2))
            cells%lon(n) = lonc*180./Pi
            cells%lat(n) = latc*180/Pi
            DO k = 0, 2
               cells%bnds_lat(k, n) = lat(k)*180./Pi
               cells%bnds_lon(k, n) = lon(k)*180./Pi
            END DO
            cells%ij(n) = d%z_up + d%iim*(j - 1) + i
         END DO
      END DO

   END SUBROUTINE set_bounds_dual

   SUBROUTINE set_bounds_edge(ind, d, cells)
      USE metric, ONLY: cell_glo
      INTEGER, INTENT(IN) :: ind
      TYPE(t_domain) :: d
      TYPE(t_cellset) :: cells
      REAL :: lon(2), lat(2)
      INTEGER :: i, j, ij, k, kk, n

      ! count edges
      n = 0
      DO j = d%jj_begin, d%jj_end
         DO i = d%ii_begin, d%ii_end
            DO k = 0, 5
               IF (d%edge_assign_domain(k, i, j) == domloc_glo_ind(ind) &
                   .AND. d%edge_assign_i(k, i, j) == i &
                   .AND. d%edge_assign_j(k, i, j) == j &
                   .AND. d%edge_assign_pos(k, i, j) == k) n = n + 1
            END DO
         END DO
      END DO
      cells%ncell = n

      IF (.NOT. ALLOCATED(cells%ij)) THEN
         ! now set bounds
         ALLOCATE (cells%ij(n), cells%lon(n), cells%lat(n), cells%ind_glo(n))
         ALLOCATE (cells%sgn(n)) ! flip sign when reading/writing
         ALLOCATE (cells%bnds_lon(2, n), cells%bnds_lat(2, n))
      END IF

      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)

      n = 0
      DO j = d%jj_begin, d%jj_end
         DO i = d%ii_begin, d%ii_end
            DO k = 0, 5
               IF (d%edge_assign_domain(k, i, j) == domloc_glo_ind(ind) &
                   .AND. d%edge_assign_i(k, i, j) == i &
                   .AND. d%edge_assign_j(k, i, j) == j &
                   .AND. d%edge_assign_pos(k, i, j) == k) THEN
                  n = n + 1
                  ij = (j - 1)*d%iim + i + u_pos(k + 1)
                  kk = MOD(k + d%delta(i, j) + 6, 6)
                  cells%ij(n) = ij
                  cells%sgn(n) = d%edge_assign_sign(k, i, j)
                  cells%ind_glo(n) = cell_glo(d%assign_cell_glo(i, j))%edge(kk) - 1
                  cells%lon(n) = lon_e(ij)*180./Pi
                  cells%lat(n) = lat_e(ij)*180./Pi

                  kk = MOD(k - 1 + 6, 6)
                  CALL xyz2lonlat(d%vertex(:, kk, i, j), lon(1), lat(1))
                  CALL xyz2lonlat(d%vertex(:, k, i, j), lon(2), lat(2))
                  cells%bnds_lon(:, n) = lon(:)*180./Pi
                  cells%bnds_lat(:, n) = lat(:)*180/Pi
               END IF
            END DO
         END DO
      END DO

   END SUBROUTINE set_bounds_edge

   SUBROUTINE set_bounds(domain_type, glo)
      TYPE(t_domain), POINTER :: domain_type(:), d
      LOGICAL :: glo
      INTEGER :: ind
      !$OMP BARRIER
      !$OMP MASTER
      ! IF glo is .TRUE. we are dealing with the global mesh, otherwise with the local mesh
      ! write_field uses the global mesh and may want halo cells
      ! output_field (XIOS) uses the local mesh and uses only own cells
      ! output_field uses edges while write_field uses only primal and dual cells
      DO ind = 1, SIZE(domain_type)
         d => domain_type(ind)
         IF (glo) THEN ! global mesh / write_field
            ! primal cell i,j is owned if d%assign_domain(i,j)==ind
            CALL set_bounds_primal(d, d%primal_own, .FALSE., d%assign_domain == ind)
            CALL set_bounds_primal(d, d%primal_all, .TRUE., d%assign_domain == ind)
            CALL set_bounds_dual(d, d%dual_own)
            CALL set_bounds_dual(d, d%dual_all)
         ELSE  ! local mesh / XIOS
            ! primal cell i,j is owned if d%own(i,j)==.TRUE.
            CALL set_bounds_primal(d, d%primal_own, .FALSE., d%own)
            CALL set_bounds_dual(d, d%dual_own)
            CALL set_bounds_edge(ind, d, d%edge_own)
         END IF
      END DO
      !$OMP END MASTER
      !$OMP BARRIER
   END SUBROUTINE set_bounds

END MODULE set_bounds_mod
