!----------------------------------------------------------------------
!     Copyright (c) 2020 Shin-ichi Takehiro. All rights reserved.
!----------------------------------------------------------------------
!
!表題  wb_spectrum_module テストプログラム :: スペクトル関数のテスト
!
!履歴  2020/08/02  竹広真一
!
program wb_spectrum_mpi_module
  use mpi
  use dc_message, only : MessageNotify
  use dc_test, only : AssertEqual
  use wb_mpi_module
  implicit none

  integer, parameter :: im=64, jm=32, nm=21, km=2, npv=2

  real(8), allocatable ::  wb_data(:,:)
  real(8), allocatable ::  pva_data(:,:,:)
  real(8), allocatable ::  b_data(:)
  real(8), allocatable ::  a_data(:)
  real(8), allocatable ::  nmb_data(:,:,:)
  real(8), allocatable ::  nb_data(:,:)


  ! 判定誤差設定
  integer, parameter :: check_digits = 10
  integer, parameter :: ignore = -11

  integer :: n,m
  integer :: ip, np, ierr

 !---------------- MPI スタート ---------------------
  call MPI_INIT(IERR)
  call MPI_COMM_RANK(MPI_COMM_WORLD,IP,IERR)
  call MPI_COMM_SIZE(MPI_COMM_WORLD,NP,IERR)

  call wb_mpi_Initial( nm, im, jm, km, npv=npv )

  allocate(wb_data(nc,kc))
  allocate(pva_data(ic,1:jc,km))
  allocate(b_data(kc))
  allocate(a_data(km))
  allocate(nmb_data(0:nm,-nm:nm,kc))
  allocate(nb_data(0:nm,kc))

  call MessageNotify('M','wb_spectrum_mpi_module_test',&
                         'wb_spectrum_module functions tests.')

  ! 2Y_1^0+3Y_2^1
  pva_data(:,:,1) = 2*sqrt(3.0D0)*sin(pv_Lat) &
       + 3*sqrt(5.0D0/6)*3.0*sin(pv_Lat)*cos(pv_Lat) * cos(pv_Lon)!
  ! 3*Y_1^0 + 2Y_1^1 + 2Y_2^0 +Y_2^-2
  pva_data(:,:,2) = 3*sqrt(3.0D0)*sin(pv_Lat) &
                + 2* sqrt(3.0D0/2)*cos(pv_Lat)*cos(pv_Lon) &
                + 2*sqrt(5.0D0)*(3.0/2*sin(pv_Lat)**2-1/2.0) &
                - sqrt(5.0D0/24)*3.0*cos(pv_Lat)**2*sin(2*pv_Lon)

  wb_data = wb_pva(pva_data)

  !--------- Energy spectrum ---------
  nmb_data=wb_spectrum_VMiss

  do n=0,nm
     do m=-n,n
        nmb_data(n,m,:) = 0.0D0
     enddo
  enddo

  if ( kp_kl(1) /= 0 ) then
     nmb_data(1,0,kp_kl(1)) = 4.0D0
     nmb_data(2,1,kp_kl(1)) = 2*3*(3.0D0/2)**2/2
     nmb_data(2,-1,kp_kl(1))= nmb_data(2,1,kp_kl(1))
  endif
  if ( kp_kl(2) /= 0 ) then
     nmb_data(1,0,kp_kl(2)) = 9.0D0
     nmb_data(1,1,kp_kl(2)) = 1*2*(2/2)**2/2
     nmb_data(1,-1,kp_kl(2))= nmb_data(1,1,kp_kl(2))
     nmb_data(2,0,kp_kl(2)) = 12.0D0
     nmb_data(2,-2,kp_kl(2))= 2*3*(1.0D0/2)**2/2
     nmb_data(2,2,kp_kl(2)) = nmb_data(2,-2,kp_kl(2))
  endif
  
  call AssertEqual(&
    message='nmb_EnergyFromStreamfunc_wb',                        &
    answer = nmb_EnergyFromStreamfunc_wb(wb_data),                &
    check =  nmb_data,                                            &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  nb_data = 0.0D0
  if ( kp_kl(1) /= 0 ) then
     nb_data(1,kp_kl(1)) = 4.0D0
     nb_data(2,kp_kl(1)) = 2 * 3 *(3.0D0/2)**2
  endif
  if ( kp_kl(2) /= 0 ) then
     nb_data(1,kp_kl(2)) = 9.0D0 + 1*2*(2/2)**2
     nb_data(2,kp_kl(2)) = 12.0D0 + 2*3*(1.0D0/2)**2
  endif
  
  call AssertEqual(&
    message='nb_EnergyFromStreamfunc_wb',                         &
    answer = nb_EnergyFromStreamfunc_wb(wb_data),                 &
    check =  nb_data,                                             &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  nmb_Data = nmb_EnergyFromStreamfunc_wb(wb_data)
  b_data = 0.0D0
  do n=0,nm
     do m=-n,n
        b_data = b_data + nmb_Data(n,m,:)
     enddo
  enddo
  a_data = reshape(aa_ab(reshape(b_data,(/kc,1/))),(/km/))

  call AssertEqual(&
    message='total energy by nm-spectrum',                        &
    answer = a_AvrLonLat_pva(                                     &
    0.5*(pva_GradLon_wb(wb_data)**2+pva_GradLat_wb(wb_data)**2) ),&
    check = a_data,                                               &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  b_data = sum(nb_EnergyFromStreamfunc_wb(wb_data),1)
  a_data = reshape(aa_ab(reshape(b_data,(/kc,1/))),(/km/))

  call AssertEqual(&
    message='total energy with n-spectrum',                       &
    answer = a_AvrLonLat_pva(                                     &
    0.5*(pva_GradLon_wb(wb_data)**2+pva_GradLat_wb(wb_data)**2) ),&
    check =  a_data,                                              &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  !--------- Enstrophy spectrum ---------
  nmb_data=wb_spectrum_VMiss
  do n=0,nm
     do m=-n,n
        nmb_data(n,m,:) = 0.0D0
     enddo
  enddo

  if ( kp_kl(1) /= 0 ) then
     nmb_data(1,0,kp_kl(1)) = 8.0D0
     nmb_data(2,1,kp_kl(1)) = 2**2*3**2*(3.0D0/2)**2/2
     nmb_data(2,-1,kp_kl(1)) = nmb_data(2,1,kp_kl(1))
  endif
  if ( kp_kl(2) /= 0 ) then
     nmb_data(1,0,kp_kl(2))  = 18.0D0
     nmb_data(1,1,kp_kl(2))  = 1**2 * 2**2 * (2/2)**2 / 2
     nmb_data(1,-1,kp_kl(2)) = nmb_data(1,1,kp_kl(2))
     nmb_data(2,0,kp_kl(2))  = 72.0D0
     nmb_data(2,-2,kp_kl(2)) = 2**2 * 3**2 * (1.0D0/2)**2 /2
     nmb_data(2,2,kp_kl(2))  = nmb_data(2,-2,kp_kl(2))
  endif

  call AssertEqual(&
    message='nmb_EnstrophyFromStreamfunc_wb',                     &
    answer = nmb_EnstrophyFromStreamfunc_wb(wb_data),             &
    check =  nmb_data,                                            &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  nb_data = 0.0
  if ( kp_kl(1) /= 0 ) then
     nb_data(1,kp_kl(1)) = 8.0D0
     nb_data(2,kp_kl(1)) = 2**2*3**2*(3.0D0/2)**2
  endif
    
  if ( kp_kl(2) /= 0 ) then
     nb_data(1,kp_kl(2)) = 18.0D0 + 1**2 * 2**2 * (2/2)**2
     nb_data(2,kp_kl(2)) = 72.0D0 + 2**2 * 3**2 * (1.0D0/2)**2
  endif

  call AssertEqual(&
    message='nb_EnstrophyFromStreamfunc_wb',                      &
    answer = nb_EnstrophyFromStreamfunc_wb(wb_data),              &
    check =  nb_data,                                             &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  nmb_Data = nmb_EnstrophyFromStreamfunc_wb(wb_data)
  b_data = 0.0D0
  do n=0,nm
     do m=-n,n
        b_data = b_data + nmb_Data(n,m,:)
     enddo
  enddo
  a_data = reshape(aa_ab(reshape(b_data,(/kc,1/))),(/km/))

  call AssertEqual(&
    message='total enstrophy by nm-spectrum',                     &
    answer = a_AvrLonLat_pva( 0.5*(pva_wb(wb_Lapla_wb(wb_data))**2)),&
    check = a_data,                                               &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  b_data = sum(nb_EnstrophyFromStreamfunc_wb(wb_data),1)
  call AssertEqual(&
    message='total energy with n-spectrum',                       &
    answer = a_AvrLonLat_pva( 0.5*(pva_wb(wb_Lapla_wb(wb_data))**2)), &
    check =  a_data,                                              &
    significant_digits = check_digits, ignore_digits = ignore     &
    )

  call MessageNotify('M','wb_spectrum_mpi_module_test', &
                     'wb_spectrum_module tests succeeded!')

 !------ MPIの終了 ------

  call MPI_FINALIZE(IERR)

end program wb_spectrum_mpi_module
