module cv3a_compress_mod
  implicit none
  private
  
  integer, parameter :: max_array_count = 100 ! Max number of arrays (of each type) added to a array_list
  integer, parameter, public :: COMPRESS_MODE_COMPRESS=1, COMPRESS_MODE_COPY=2 ! Constants for different compress modes
  integer, save, public :: compress_mode = COMPRESS_MODE_COMPRESS ! Current compress mode. Can be changed by user, but be sure to use only one type of compression.
  !$omp threadprivate(compress_mode)
  
  ! Containers for arrays (Internal structure)
  type array_r_t
    real, pointer, contiguous :: buf(:,:,:) => null() ! max dim for real is 3
  end type
  type array_i_t
    integer, pointer :: buf(:) => null() ! max dim for integer is 1
  end type
  
  ! List of arrays (Black box structure)
  type array_list
    integer :: array_count_r = 0 ! number of real arrays
    type(array_r_t) :: arrays_r(2,max_array_count) ! array of real arrays
    logical :: arrays_r_init(max_array_count) = .false. ! Should arrays (2,:) be initialized before uncompress?
    integer :: array_count_i = 0 ! number of int arrays
    type(array_i_t) :: arrays_i(2,max_array_count) ! array of int arrays
    logical :: arrays_i_init(max_array_count) = .false. ! Should arrays (2,:) be initialized before uncompress?
  end type
  
  ! Information initialized by compress to transmit instructions for uncompress
  type compress_data_t
    integer :: ncum ! Number of cells in compressed arrays (first dimension)
    integer, allocatable :: idcum(:) ! Correspondance between compressed and uncompressed arrays uncomp(idcum(i)) == comp(i) 
    logical, allocatable :: mask(:) ! Mask used for compression
  end type 
  
  public :: array_list, add_array_i1, add_array_r1, add_array_r2, add_array_r3, &
            cv3a_compress, cv3a_uncompress, compress_data_t
  
  contains
  
  ! Add a pair of array to an array list (integer(:))
  subroutine add_array_i1(arrays, array_in, array_out, init)
    type(array_list) :: arrays ! List of arrays to add arrays to
    integer, target :: array_in(:) ! Input array (should not be modified)
    integer, target :: array_out(:) ! Output array (will be modified)
    logical, optional, intent(in) :: init ! For uncompress : zero-initialize output array before uncompress
    integer :: c   
    arrays%array_count_i = arrays%array_count_i + 1
    c = arrays%array_count_i 
    if( c > max_array_count ) call abort_physic('cv3a_compress', "Too many arrays : increase cv3a_compress_mod%max_array_count",1)
    if( present(init) ) arrays%arrays_i_init(c)=init
    arrays%arrays_i(1,c)%buf => array_in
    arrays%arrays_i(2,c)%buf => array_out
  end subroutine

  ! Add a pair of array to an array list (real(:))
  subroutine add_array_r1(arrays, array_in, array_out, init)
    type(array_list) :: arrays
    real, target :: array_in(:), array_out(:)
    logical, optional, intent(in) :: init
    integer :: c
    arrays%array_count_r = arrays%array_count_r + 1
    c = arrays%array_count_r 
    if( c > max_array_count ) call abort_physic('cv3a_compress', "Too many arrays : increase cv3a_compress_mod%max_array_count",1)
    if( present(init) ) arrays%arrays_r_init(c)=init 
    arrays%arrays_r(1,c)%buf(1:size(array_in),1:1,1:1) => array_in
    arrays%arrays_r(2,c)%buf(1:size(array_out),1:1,1:1) => array_out
  end subroutine
  
  ! Add a pair of array to an array list (real(:,:))
  subroutine add_array_r2(arrays, array_in, array_out, init)
    type(array_list) :: arrays
    real, target, contiguous :: array_in(:,:), array_out(:,:)
    logical, optional, intent(in) :: init
    integer :: c
    arrays%array_count_r = arrays%array_count_r + 1
    c = arrays%array_count_r 
    if( c > max_array_count ) call abort_physic('cv3a_compress', "Too many arrays : increase cv3a_compress_mod%max_array_count",1)
    if( present(init) ) arrays%arrays_r_init(c)=init   
    arrays%arrays_r(1,c)%buf(1:size(array_in,1),1:size(array_in,2),1:1) => array_in
    arrays%arrays_r(2,c)%buf(1:size(array_out,1),1:size(array_out,2),1:1) => array_out
  end subroutine
  
  ! Add a pair of array to an array list (real(:,:,:))
  subroutine add_array_r3(arrays, array_in, array_out, init)
    type(array_list) :: arrays
    real, target :: array_in(:,:,:), array_out(:,:,:)
    logical, optional, intent(in) :: init
    integer :: c
    arrays%array_count_r = arrays%array_count_r + 1
    c = arrays%array_count_r 
    if( c > max_array_count ) call abort_physic('cv3a_compress', "Too many arrays : increase cv3a_compress_mod%max_array_count",1) 
    if( present(init) ) arrays%arrays_r_init(c)=init   
    arrays%arrays_r(1,c)%buf => array_in
    arrays%arrays_r(2,c)%buf => array_out
  end subroutine 
    
  ! Compress arrays from 'arrays' according to 'mask'.
  ! Input arrays in 'arrays' should have 'len' as first dimension
  ! 2D arrays (real only) are compressed on every vertical layer
  ! Note : compressing 3d arrays is not supported yet 
  subroutine cv3a_compress(len, mask, arrays, d)
    integer, INTENT (IN) :: len ! lenght of the contiguous dimension of arrays to compress
    logical, INTENT (IN) :: mask(len) ! Mask of convective cells
    type(array_list), intent(inout) :: arrays ! List of arrays to compress (first compressed into second)
    type(compress_data_t), intent(out) :: d ! Used to store data about compression to be able to uncompress
    
    !real, parameter :: comp_threshold = 1
    integer :: i, ncum, ncum_max
    
    ! Output arrays are supposed to be able to contain compressed data
    ncum_max = max( size(arrays%arrays_i(2,1)%buf), len )  
    allocate(d%idcum(ncum_max))
    allocate(d%mask(len))
    d%mask(:) = mask(:)
    
    ! Generate permutation idcum() and number of active elements 'ncum' for compression
    ncum = 0
    do i = 1, len
      if (mask(i)) then
        ncum = ncum + 1
        d%idcum(ncum) = i
      end if
    end do
    if( ncum > ncum_max ) call abort_physic("cv3a_compress", "internal error : idcum too small", 1)
    
    !if (ncum < len * comp_threshold) compress_mode=COMPRESS_MODE_COMPRESS    
    select case (compress_mode)
      case (COMPRESS_MODE_COMPRESS)
        call cv3a_compress_compress(arrays, ncum, d%idcum)
      case (COMPRESS_MODE_COPY)
        call cv3a_compress_copy(len, arrays, ncum, d%idcum)
      case default
        call abort_physic("cv3a_compress", "Unknown compress mode", 1)
    end select
    
    !if( ncum /= get_compress_size(len, mask) ) call abort_physic("cv3a_compress", "internal error : get_compress_size mismatch", 1)
    
    d%ncum = ncum

  end subroutine

  ! cv3a_compress in mode COMPRESS_MODE_COMPRESS
  ! Output arrays have ncum elements that correspond to the idcum(i) 
  ! elements of the input arrays : a_out(i) = a_in(idcum(i))
  ! compressing 3D arrays is not supported
  subroutine cv3a_compress_compress(arrays, ncum, idcum)
    type(array_list), intent(inout) :: arrays ! List of arrays to compress (first compressed into second)
    integer, intent(in) :: ncum ! Number of cells in compressed arrays (first dimension)
    integer, intent(in) :: idcum(ncum) ! Correspondance between compressed and uncompressed arrays uncomp(idcum(i)) == comp(i) 
    
    integer :: i,j,k,a
     
    ! Compress int arrays
    do a = 1, arrays%array_count_i
      do i = 1, ncum
        arrays%arrays_i(2,a)%buf(i) = arrays%arrays_i(1,a)%buf(idcum(i))
      end do
    end do
    ! Compress real arrays
    do a = 1, arrays%array_count_r
      if( size(arrays%arrays_r(2,a)%buf, 3) /= 1 ) call abort_physic("cv3a_compress", "cannot compress 3D array", 1)
      do k = 1, size(arrays%arrays_r(2,a)%buf, 3)
        do j = 1, size(arrays%arrays_r(2,a)%buf, 2) 
         do i = 1, ncum
           arrays%arrays_r(2,a)%buf(i,j,k) = arrays%arrays_r(1,a)%buf(idcum(i),j,k)
         end do
        end do
      end do            
    end do  
  end subroutine
  
  ! cv3a_compress in mode COMPRESS_MODE_COPY
  subroutine cv3a_compress_copy(len, arrays, ncum, idcum)
    integer, INTENT (IN) :: len ! lenght of the contiguous dimension of arrays to compress
    type(array_list), intent(inout) :: arrays ! List of arrays to compress (first compressed into second)
    integer, intent(out) :: ncum ! Number of cells in compressed arrays (first dimension)
    integer, intent(out) :: idcum(len) ! Correspondance between compressed and uncompressed arrays uncomp(idcum(i)) == comp(i) 
    integer :: i  

    ! Ignore and ovewrite computed permutation
    ncum = len
    DO i = 1,len
      idcum(i) = i
    ENDDO  
    ! Copy int arrays
    do i = 1, arrays%array_count_i
      arrays%arrays_i(2,i)%buf(:) = arrays%arrays_i(1,i)%buf(:)
    end do
    ! Copy real arrays
    do i = 1, arrays%array_count_r
      arrays%arrays_r(2,i)%buf(:,:,:) = arrays%arrays_r(1,i)%buf(:,:,:)
    end do
  end subroutine
  
  ! Arrays in 'arrays' are uncompressed, 
  ! non-convective cells are zero-initialized if init=.true. was used when adding the array 
  ! 2D arrays are uncompressed up to level 'nl' (cv3param.h)
  ! 3D arrays are uncompressed up to level (:,nl,nl) (cv3param.h)
  ! after those levels cells are zero-initialized when init=.true.
  subroutine cv3a_uncompress(len, d, arrays) 
    integer, intent(in) :: len ! Number of cells in uncompressed arrays (first dimension)
    type(compress_data_t), intent(in) :: d ! Data to uncompress arrays
    type(array_list), intent(inout) :: arrays ! List of arrays to compress (first compressed into second)
    
    select case (compress_mode)
      case (COMPRESS_MODE_COMPRESS)
        call cv3a_uncompress_uncompress(d%ncum, d%idcum, arrays)
      case (COMPRESS_MODE_COPY)
        call cv3a_uncompress_copy(d%ncum, d%mask, arrays)
      case default
        call abort_physic("cv3a_uncompress", "Unknown uncompress mode", 1)
    end select  
  end subroutine
  
  ! cv3a_uncompress in mode COMPRESS_MODE_COMPRESS
  ! a_out(idcum(i)) = a_in(i)
  subroutine cv3a_uncompress_uncompress(ncum, idcum, arrays) 
    integer, intent(in) :: ncum ! Number of cells in compressed arrays (first dimension)
    integer, intent(in) :: idcum(ncum) ! Correspondance between compressed and uncompressed arrays uncomp(idcum(i)) == comp(i) 
    type(array_list), intent(inout) :: arrays ! List of arrays to compress (first compressed into second)
    
    include "cv3param.h"
    
    integer :: a, k, j, i  
          
    ! uncompress int arrays
    do a = 1, arrays%array_count_i
        if(arrays%arrays_i_init(a)) arrays%arrays_i(2,a)%buf(:) = 0
        do i = 1, ncum
          arrays%arrays_i(2,a)%buf(idcum(i)) = arrays%arrays_i(1,a)%buf(i)
        end do
     end do
    ! Compress real arrays
    do a = 1, arrays%array_count_r
      if(arrays%arrays_r_init(a))  arrays%arrays_r(2,a)%buf(:,:,:) = 0
      do k = 1, min(size(arrays%arrays_r(2,a)%buf, 3), nl)
        do j = 1, min(size(arrays%arrays_r(2,a)%buf, 2), nl)
          do i = 1, ncum
            arrays%arrays_r(2,a)%buf(idcum(i),j,k) = arrays%arrays_r(1,a)%buf(i,j,k)
          end do
        end do
      end do
    end do               
  end subroutine

  subroutine cv3a_uncompress_copy(len, mask, arrays)
    integer, intent(in) :: len
    logical, intent(in) :: mask(:)
    type(array_list), intent(inout) :: arrays ! List of arrays to compress (first compressed into second)
    
    include "cv3param.h"
    
    integer :: a, i, j, k
    
    ! Copy int arrays
    do a = 1, arrays%array_count_i
      if(arrays%arrays_i_init(a)) arrays%arrays_i(2,a)%buf(:) = 0
      do i = 1, len
        if(mask(i)) arrays%arrays_i(2,a)%buf(i) = arrays%arrays_i(1,a)%buf(i)
      end do
    end do
    ! Copy real arrays
    do a = 1, arrays%array_count_r
      if(arrays%arrays_i_init(a)) arrays%arrays_r(2,a)%buf(:,:,:) = 0        
      do k = 1, min(size(arrays%arrays_r(2,a)%buf, 3), nl)
        do j = 1, min(size(arrays%arrays_r(2,a)%buf, 2), nl)
          do i = 1, len
            if(mask(i)) arrays%arrays_r(2,a)%buf(:,j,k) = merge( arrays%arrays_r(1,a)%buf(:,j,k), 0., mask)
          end do
        end do
      end do
    end do
  end subroutine
  
end module