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
       ENDDO
    ENDDO
    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
       ENDDO
    ENDDO
    cells%ncell = n

    IF ( .NOT. ALLOCATED( cells%ind_glo ) ) THEN 
        ! now set bounds
        ALLOCATE(cells%ind_glo(n)) ! not set but must be 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
       ENDDO
    ENDDO
    
    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
       ENDDO
    ENDDO

  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