MODULE field_mod
  USE genmod
  IMPLICIT NONE

  INTEGER,PARAMETER :: field_T=1
  INTEGER,PARAMETER :: field_U=2
  INTEGER,PARAMETER :: field_Z=3

  INTEGER,PARAMETER :: type_integer=1
  INTEGER,PARAMETER :: type_real=2
  INTEGER,PARAMETER :: type_logical=3

  TYPE t_field
     CHARACTER(30)      :: name
     LOGICAL :: ondevice !< flag if field is allocated on device as well
     INTEGER :: ndim
     INTEGER :: field_type
     INTEGER :: data_type
     INTEGER :: dim3
     INTEGER :: dim4
     REAL(rstd), POINTER, CONTIGUOUS :: rval2d(:) => NULL()
     REAL(rstd), POINTER, CONTIGUOUS :: rval3d(:,:) => NULL()
     REAL(rstd), POINTER, CONTIGUOUS :: rval4d(:,:,:) => NULL()
     INTEGER, POINTER, CONTIGUOUS :: ival2d(:) => NULL()
     INTEGER, POINTER, CONTIGUOUS :: ival3d(:,:) => NULL()
     INTEGER, POINTER, CONTIGUOUS :: ival4d(:,:,:) => NULL()
     LOGICAL, POINTER, CONTIGUOUS :: lval2d(:) => NULL()
     LOGICAL, POINTER, CONTIGUOUS :: lval3d(:,:) => NULL()
     LOGICAL, POINTER, CONTIGUOUS :: lval4d(:,:,:) => NULL()
  END TYPE t_field

  INTERFACE get_val
     MODULE PROCEDURE getval_r2d
     MODULE PROCEDURE getval_r3d
     MODULE PROCEDURE getval_r4d
     MODULE PROCEDURE getval_i2d
     MODULE PROCEDURE getval_i3d
     MODULE PROCEDURE getval_i4d
     MODULE PROCEDURE getval_l2d
     MODULE PROCEDURE getval_l3d
     MODULE PROCEDURE getval_l4d
  END INTERFACE

  INTERFACE ASSIGNMENT(=)
     MODULE PROCEDURE getval_r2d
     MODULE PROCEDURE getval_r3d
     MODULE PROCEDURE getval_r4d
     MODULE PROCEDURE getval_i2d
     MODULE PROCEDURE getval_i3d
     MODULE PROCEDURE getval_i4d
     MODULE PROCEDURE getval_l2d
     MODULE PROCEDURE getval_l3d
     MODULE PROCEDURE getval_l4d
  END INTERFACE

  PRIVATE :: allocate_field_, deallocate_field_

CONTAINS

  !====================================== allocate_field ===================================

  SUBROUTINE allocate_field_glo(field,field_type,data_type,dim1,dim2,name)
    USE domain_mod
    TYPE(t_field),POINTER :: field(:)
    INTEGER,INTENT(IN) :: field_type
    INTEGER,INTENT(IN) :: data_type
    INTEGER,OPTIONAL :: dim1,dim2
    CHARACTER(*), OPTIONAL :: name
    INTEGER :: ind

    ALLOCATE(field(ndomain_glo))
    DO ind=1,ndomain_glo
       CALL allocate_field_(domain_glo(ind), field(ind), field_type, data_type, dim1, dim2, name)
    ENDDO

  END SUBROUTINE allocate_field_glo

  SUBROUTINE allocate_field(field, field_type, data_type, dim3, dim4, name, ondevice)
    USE domain_mod
    USE omp_para
    TYPE(t_field), POINTER :: field(:)
    INTEGER, INTENT(IN)    :: field_type
    INTEGER, INTENT(IN)    :: data_type
    INTEGER, OPTIONAL      :: dim3,dim4
    CHARACTER(*), OPTIONAL :: name
    LOGICAL, INTENT(IN), OPTIONAL :: ondevice
    INTEGER :: ind
    !$OMP BARRIER
    !$OMP MASTER
    ALLOCATE(field(ndomain))
    !$OMP END MASTER
    !$OMP BARRIER

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master) CYCLE
       CALL allocate_field_(domain(ind), field(ind), field_type, data_type, dim3, dim4, name, ondevice)
    END DO
    !$OMP BARRIER

  END SUBROUTINE allocate_field

  SUBROUTINE allocate_fields(nfield,field,field_type,data_type,dim3,dim4,name, ondevice)
    USE domain_mod
    USE omp_para
    INTEGER, INTENT(IN)     :: nfield
    TYPE(t_field), POINTER  :: field(:,:)
    INTEGER, INTENT(IN)     :: field_type
    INTEGER, INTENT(IN)     :: data_type
    INTEGER, OPTIONAL       :: dim3,dim4
    CHARACTER(*), OPTIONAL  :: name
    LOGICAL, INTENT(IN), OPTIONAL :: ondevice
    INTEGER :: i, ind
    !$OMP BARRIER
    !$OMP MASTER
    ALLOCATE(field(ndomain,nfield))
    !$OMP END MASTER
    !$OMP BARRIER
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master) CYCLE
       DO i=1,nfield
          CALL allocate_field_(domain(ind), field(ind,i),field_type, data_type, dim3, dim4, name, ondevice)
       END DO
    END DO
    !$OMP BARRIER

  END SUBROUTINE allocate_fields

  SUBROUTINE allocate_field_(dom, field, field_type, data_type, dim3, dim4, name, ondevice)
    USE domain_mod
    USE omp_para
    TYPE(t_domain)         :: dom
    TYPE(t_field)          :: field
    INTEGER, INTENT(IN)    :: field_type
    INTEGER, INTENT(IN)    :: data_type
    INTEGER, OPTIONAL      :: dim3,dim4
    CHARACTER(*), OPTIONAL :: name
    LOGICAL, INTENT(IN), OPTIONAL :: ondevice

    INTEGER :: ij_size

    IF(PRESENT(name)) THEN
       field%name = name
    ELSE
       field%name = '(undefined)'
    END IF

    IF (PRESENT(dim4)) THEN
       field%ndim=4
       field%dim4=dim4
       field%dim3=dim3
    ELSE IF (PRESENT(dim3)) THEN
       field%ndim=3
       field%dim3=dim3
       field%dim4=1
    ELSE
       field%ndim=2
       field%dim3=1
       field%dim4=1
    ENDIF


    field%data_type=data_type
    field%field_type=field_type

    IF (field_type==field_T) THEN
       ij_size=dom%iim*dom%jjm
    ELSE IF (field_type==field_U) THEN
       ij_size=3*dom%iim*dom%jjm
    ELSE IF (field_type==field_Z) THEN
       ij_size=2*dom%iim*dom%jjm
    ENDIF

    IF (data_type==type_integer) ALLOCATE(field%ival4d(ij_size, field%dim3, field%dim4))
    IF (data_type==type_real)    ALLOCATE(field%rval4d(ij_size, field%dim3, field%dim4))
    IF (data_type==type_logical) ALLOCATE(field%lval4d(ij_size, field%dim3, field%dim4))

    IF (field%ndim==3) THEN
       IF (data_type==type_integer) field%ival3d => field%ival4d(:,:,1)
       IF (data_type==type_real)    field%rval3d => field%rval4d(:,:,1)
       IF (data_type==type_logical) field%lval3d => field%lval4d(:,:,1)

    ELSE IF (field%ndim==2) THEN
       IF (data_type==type_integer) field%ival2d => field%ival4d(:,1,1)
       IF (data_type==type_real)    field%rval2d => field%rval4d(:,1,1)
       IF (data_type==type_logical) field%lval2d => field%lval4d(:,1,1)

    ENDIF

    field%ondevice = .FALSE.
    IF (PRESENT(ondevice)) THEN
       IF (ondevice) CALL create_device_field(field)
    END IF

  END SUBROUTINE allocate_field_

  !==================================== deallocate_field ===================================

  SUBROUTINE deallocate_field_glo(field)
    USE domain_mod
    TYPE(t_field),POINTER :: field(:)
    INTEGER :: ind
    DO ind=1,ndomain_glo
       CALL deallocate_field_(field(ind))
    END DO
    DEALLOCATE(field)
  END SUBROUTINE deallocate_field_glo

  SUBROUTINE deallocate_field(field)
    USE domain_mod
    USE omp_para
    TYPE(t_field),POINTER :: field(:)
    INTEGER :: ind
    !$OMP BARRIER
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master) CYCLE
       CALL deallocate_field_(field(ind))
    END DO
    !$OMP BARRIER
    !$OMP MASTER
    DEALLOCATE(field)
    !$OMP END MASTER
    !$OMP BARRIER
  END SUBROUTINE deallocate_field

  SUBROUTINE deallocate_fields(field)
    USE domain_mod
    USE omp_para
    TYPE(t_field),POINTER :: field(:,:)
    INTEGER :: i, ind
    !$OMP BARRIER
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master) CYCLE
       DO i=1,SIZE(field,2)
          CALL deallocate_field_(field(ind,i))
       END DO
    END DO
    !$OMP BARRIER
    !$OMP MASTER
    DEALLOCATE(field)
    !$OMP END MASTER
    !$OMP BARRIER
  END SUBROUTINE deallocate_fields

  SUBROUTINE deallocate_field_(field)
    USE domain_mod
    USE omp_para
    TYPE(t_field) :: field
    INTEGER :: data_type
    data_type=field%data_type
    IF (data_type==type_real) THEN
       IF (field%ondevice) THEN
          !$acc exit data delete(field%rval4d(:,:,:))
          CONTINUE
       END IF
       DEALLOCATE(field%rval4d)
    END IF
    IF (data_type==type_integer) THEN
       IF (field%ondevice) THEN
          !$acc exit data delete(field%ival4d(:,:,:))
          CONTINUE
       END IF
       DEALLOCATE(field%ival4d)
    END IF
    IF (data_type==type_logical) THEN
       IF (field%ondevice) THEN
          !$acc exit data delete(field%lval4d(:,:,:))
          CONTINUE
       END IF
       DEALLOCATE(field%lval4d)
    END IF

  END SUBROUTINE deallocate_field_

  !====================================== getval ===================================

  SUBROUTINE getval_r2d(field_pt,field)
    REAL(rstd), POINTER, INTENT(INOUT) :: field_pt(:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=2 .OR. field%data_type/=type_real) THEN
       PRINT *, 'getval_r2d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%rval2d
  END SUBROUTINE getval_r2d

  SUBROUTINE getval_r3d(field_pt,field)
    REAL(rstd), POINTER, INTENT(INOUT) :: field_pt(:,:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=3 .OR. field%data_type/=type_real) THEN
       PRINT *, 'getval_r3d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%rval3d
  END SUBROUTINE getval_r3d

  SUBROUTINE getval_r4d(field_pt,field)
    REAL(rstd), POINTER, INTENT(INOUT) :: field_pt(:,:,:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=4 .OR. field%data_type/=type_real) THEN
       PRINT *, 'getval_r4d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%rval4d
  END SUBROUTINE getval_r4d

  SUBROUTINE getval_i2d(field_pt,field)
    INTEGER, POINTER, INTENT(INOUT) :: field_pt(:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=2 .OR. field%data_type/=type_integer) THEN
       PRINT *, 'getval_i2d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%ival2d
  END SUBROUTINE getval_i2d

  SUBROUTINE getval_i3d(field_pt,field)
    INTEGER, POINTER, INTENT(INOUT) :: field_pt(:,:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=3 .OR. field%data_type/=type_integer) THEN
       PRINT *, 'getval_i3d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%ival3d
  END SUBROUTINE getval_i3d

  SUBROUTINE getval_i4d(field_pt,field)
    INTEGER, POINTER, INTENT(INOUT) :: field_pt(:,:,:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=4 .OR. field%data_type/=type_integer) THEN
       PRINT *, 'getval_i4d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%ival4d
  END SUBROUTINE getval_i4d

  SUBROUTINE getval_l2d(field_pt,field)
    LOGICAL, POINTER, INTENT(INOUT) :: field_pt(:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=2 .OR. field%data_type/=type_logical) THEN
       PRINT *, 'getval_l2d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%lval2d
  END SUBROUTINE getval_l2d

  SUBROUTINE getval_l3d(field_pt,field)
    LOGICAL, POINTER, INTENT(INOUT) :: field_pt(:,:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=3 .OR. field%data_type/=type_logical) THEN
       PRINT *, 'getval_l3d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%lval3d
  END SUBROUTINE getval_l3d

  SUBROUTINE getval_l4d(field_pt,field)
    LOGICAL, POINTER, INTENT(INOUT) :: field_pt(:,:,:)
    TYPE(t_field),INTENT(IN) :: field

    IF (field%ndim/=4 .OR. field%data_type/=type_logical) THEN
       PRINT *, 'getval_l4d : bad pointer assignment with ' // TRIM(field%name)
       STOP
    END IF
    field_pt=>field%lval4d
  END SUBROUTINE getval_l4d

  !===================== Data transfer between host (CPU) and device (GPU) =========================

  SUBROUTINE update_device_field(field)
    USE domain_mod
    USE omp_para
    TYPE(t_field) :: field(:)
    INTEGER :: ind

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master) CYCLE

       IF (.NOT. field(ind)%ondevice) CALL create_device_field(field(ind))
       IF (field(ind)%data_type==type_real) THEN
          !$acc update device(field(ind)%rval4d(:,:,:)) async
          CONTINUE
       END IF
       IF (field(ind)%data_type==type_integer) THEN
          !$acc update device(field(ind)%ival4d(:,:,:)) async
          CONTINUE
       END IF
       IF (field(ind)%data_type==type_logical) THEN
          !$acc update device(field(ind)%lval4d(:,:,:)) async
          CONTINUE
       END IF

    ENDDO
    !$OMP BARRIER
  END SUBROUTINE update_device_field

  SUBROUTINE update_host_field(field)
    USE domain_mod
    USE omp_para
    TYPE(t_field) :: field(:)
    INTEGER :: ind

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master) CYCLE

       IF (field(ind)%ondevice) THEN
          IF (field(ind)%data_type==type_real) THEN
             !$acc update host(field(ind)%rval4d(:,:,:)) async
             CONTINUE
          END IF
          IF (field(ind)%data_type==type_integer) THEN
             !$acc update host(field(ind)%ival4d(:,:,:)) async
             CONTINUE
          END IF
          IF (field(ind)%data_type==type_logical) THEN
             !$acc update host(field(ind)%lval4d(:,:,:)) async
             CONTINUE
          END IF

       END IF
    ENDDO
    !$acc wait
    !$OMP BARRIER
  END SUBROUTINE update_host_field

  SUBROUTINE create_device_field(field)
    TYPE(t_field) :: field

    IF (field%ondevice) THEN
       PRINT *, "Field is already on device !"
       STOP 1
    END IF
    IF (field%data_type==type_real) THEN
       !$acc enter data create(field%rval4d(:,:,:)) async
    END IF
    IF (field%data_type==type_integer) THEN
       !$acc enter data create(field%ival4d(:,:,:)) async
    END IF
    IF (field%data_type==type_logical) THEN
       !$acc enter data create(field%lval4d(:,:,:)) async
    END IF

    field%ondevice = .TRUE.
  END SUBROUTINE create_device_field

END MODULE field_mod