!===============================================================================
!
! cg_init_modules.F90
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2012-2016 Hinnerk Stueben
!
! This file is part of BQCD -- Berlin Quantum ChromoDynamics program
!
! BQCD is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! BQCD is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with BQCD.  If not, see <http://www.gnu.org/licenses/>.
!
!-------------------------------------------------------------------------------
# define LANGUAGE_FORTRAN
# include "defs.h"
# include "simd.h"
# include "alignment.h"

!-------------------------------------------------------------------------------
subroutine cg_init_modules()

  implicit none
  logical, save :: initialised = .false.

  if (initialised) return

  call init_module_cg_vol()
  call init_module_cg_nn()
  call init_module_cg_surface()
  call init_module_cg_arrays()
  call init_module_cg_conf()
  call init_module_cg_procs()
  call init_module_cg_proj_xbound()
  call init_module_cg_threads()

  initialised = .true.

end

!-------------------------------------------------------------------------------
subroutine init_module_cg_vol()

  use module_lattice
  use module_cg_vol
  use module_vol
  implicit none

  CHECK(mod(nxh, SIMD_WIDTH_R4) == 0)
  CHECK(mod(nxh, SIMD_WIDTH_R8) == 0)
  CHECK(mod(volh, SIMD_WIDTH_R4) == 0)
  CHECK(mod(volh, SIMD_WIDTH_R8) == 0)
  CHECK(mod(volh_tot, SIMD_WIDTH_R4) == 0)
  CHECK(mod(volh_tot, SIMD_WIDTH_R8) == 0)
  
  volh_r4  = volh
  volh_r4s = volh / SIMD_WIDTH_R4
  volh_r8  = volh
  volh_r8s = volh / SIMD_WIDTH_R8

  volh_tot_r4  = volh_tot
  volh_tot_r4s = volh_tot / SIMD_WIDTH_R4
  volh_tot_r8  = volh_tot
  volh_tot_r8s = volh_tot / SIMD_WIDTH_R8

end

!-------------------------------------------------------------------------------
subroutine init_module_cg_nn()

  use module_cg_nn
  use module_cg_vol
  use module_lattice
  use module_nn
  use module_vol
  use module_function_decl
  implicit none

  integer :: fb, mu, eo, j, i, ii, nn_j, nn_fwd, first, last
  integer :: dir, x, y, z, t, xyzt(DIM)
  integer :: i_simd, i_simd_old(EVEN:ODD), neighbour, neighbour_simd
  integer, external :: xyzt2i, e_o, i_periodic

FOR _rr in _r4 _r8
SET _rrs = _rr %% s
SET SIMD_WIDTH = SIMD_WIDTH %% $(toupper _rr)

  nn %% _rr => nn
  allocate(nn %% _rrs(volh_tot %% _rrs, EVEN:ODD, DIM, FWD:BWD))

  do fb = FWD, BWD
     if (fb == FWD) then
        dir = +1
     else
        dir = -1
     endif

     do mu = 1, DIM
        i_simd_old(:) = 0

        do t = 0, nt - 1
        do z = 0, nz - 1
        do y = 0, ny - 1
        do x = 0, nx - 1
           xyzt = (/x,y,z,t/)

           i = xyzt2i(xyzt)

           if (mod(i - 1, SIMD_WIDTH) == 0) then
              eo = e_o(xyzt)
              neighbour = nn(i, eo, mu, fb)

              i_simd         = 1 + (i         - 1) / SIMD_WIDTH
              neighbour_simd = 1 + (neighbour - 1) / SIMD_WIDTH

              if (mu == 1) then
              if (neighbour_simd == i_simd) then
              if (neighbour /= i) then
                 if (SIMD_WIDTH == nxh) then 
                    neighbour_simd = -neighbour_simd
                 else
                    xyzt = (/x,y,z,t/)
                    xyzt(mu) = i_periodic(xyzt(mu) + dir * 2 * SIMD_WIDTH, N(mu))
                    neighbour_simd = 1 + (xyzt2i(xyzt) - 1) / SIMD_WIDTH
                 endif
              endif
              endif
              endif

              nn%%_rrs(i_simd, eo, mu, fb) = neighbour_simd
              
              ASSERT(i_simd == i_simd_old(eo) + 1)
              i_simd_old(eo) = i_simd
           endif

        enddo
        enddo
        enddo
        enddo

        ASSERT(i_simd_old(EVEN) == volh%%_rrs)
        ASSERT(i_simd_old(ODD) == volh%%_rrs)

     enddo
  enddo

ENDSET
ENDSET
ENDFOR

end

!-------------------------------------------------------------------------------
subroutine init_module_cg_surface()

  use module_surface
  use module_cg_surface
  implicit none
  integer :: mu, fb

  surface2_r4 => surface2
  surface2_r8 => surface2
  surface_halo2_r4 => surface_halo2
  surface_halo2_r8 => surface_halo2

  do mu = 2, DIM
  FOR surface in surface2 surface_halo2
  FOR _rr in _r4 _r8
  SET _rrs = _rr %% s
  SET SIMD_WIDTH = SIMD_WIDTH %% $(toupper _rr)

     CHECK(mod(surface(mu)%i_start(FWD) - 1, SIMD_WIDTH) == 0)  
     CHECK(mod(surface(mu)%i_start(BWD) - 1, SIMD_WIDTH) == 0)  
     CHECK(mod(surface(mu)%blocklength, SIMD_WIDTH) == 0)  
     CHECK(mod(surface(mu)%stride, SIMD_WIDTH) == 0)  

     surface%%_rrs(mu)%i_start(FWD) = 1 + (surface(mu)%i_start(FWD) - 1) / SIMD_WIDTH
     surface%%_rrs(mu)%i_start(BWD) = 1 + (surface(mu)%i_start(BWD) - 1) / SIMD_WIDTH
     surface%%_rrs(mu)%count        = surface(mu)%count
     surface%%_rrs(mu)%blocklength  = surface(mu)%blocklength / SIMD_WIDTH
     surface%%_rrs(mu)%stride       = surface(mu)%stride / SIMD_WIDTH
     surface%%_rrs(mu)%nsites       = surface(mu)%nsites
  ENDSET
  ENDSET
  ENDFOR
  ENDFOR
  enddo
end

!-------------------------------------------------------------------------------
subroutine init_module_cg_procs()

  use module_cg_procs
  use module_input
  use module_function_decl
  implicit none
  character(8) :: r, rs, sc, g, c, d

  select case (input%tuning_cg_precision)
  case (32); r = "_r4"
  case (64); r = "_r8"
  case default; call die("init_module_cg_procs(): unknown tuning_cg_precision")
  end select

  select case (input%tuning_cg_spincol)
  case (1); sc = "_sc"
  case (22); sc = "_sc22"
  case default; call die("init_module_cg_procs(): unknown tuning_cg_spincol")
  end select

  select case (input%tuning_cg_clover)
  case (1); c = "_c1"
  case (2); c = "_c2"
  case default; call die("init_module_cg_procs(): unknown tuning_cg_clover")
  end select

  select case (input%tuning_cg_d)
  case (2);   d = "d2";    g = "_g1"
  case (21);  d = "d21";   g = "_g1"
  case (25);  d = "d25";   g = "_g1"; CHECK(input%processes(1) == 1)
  case (35);  d = "d35";   g = "_g1"; CHECK(input%processes(1) == 1); CHECK(num_threads() >= 2)
  case (102); d = "d102";  g = "_g2"
  case default; call die("init_module_cg_procs(): unknown tuning_cg_d")
  end select

  select case (input%tuning_cg_simd)
  case (0); rs = r
  case (1); rs = trim(r) // "s"
            CHECK(input%processes(1) == 1)
  case default; call die("init_module_cg_procs(): unknown tuning_cg_simd")
  end select


  call cg_set_proc(cg_procs%d,     "cg_" // trim(d)           // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%d_dag, "cg_" // trim(d) // "_dag" // trim(sc) // trim(rs))
  
  call cg_set_proc(cg_procs%co,    "cg_co"    // trim(c) // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%c2,    "cg_c2"    // trim(c) // trim(sc) // trim(rs))

  call cg_set_proc(cg_procs%loop1, "cg_loop1" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%loop2, "cg_loop2" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%loop3, "cg_loop3" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%axpy,  "cg_axpy"  // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%xpby,  "cg_xpby"  // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%norm2, "cg_norm2" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%cdotc, "cg_cdotc" // trim(sc) // trim(rs))

  call cg_set_proc(cg_procs%copyin_spincol,  "cg_copyin"  // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%copyout_spincol, "cg_copyout" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%copy_caxpy,      "cg_copy_caxpy" // trim(sc) // trim(rs))

  call cg_set_proc(cg_procs%copyin_gauge,    "cg_copyin"  // trim(g) // trim(rs))
  call cg_set_proc(cg_procs%copyin_clover,   "cg_copyin"  // trim(c) // trim(rs))

  call cg_set_proc(cg_procs%kernel,          "cg_kernel2" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%kernel_cg,       "cg_kernel_cg" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%kernel_mcg,      "cg_kernel_mcg" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%kernel_mbcg,     "cg_kernel_mbcg" // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%mm,              "cg_mm"      // trim(sc) // trim(rs))
  call cg_set_proc(cg_procs%ww,              "cg_ww"      // trim(sc) // trim(rs))

end

!-------------------------------------------------------------------------------
subroutine init_module_cg_arrays()

  use module_cg_arrays
  use module_input
  use module_vol
  implicit none
  integer :: mu, fb

  real(8), dimension(:), allocatable :: x
  integer :: n
  ATTRIBUTE_ALIGN(x)

  n = SIZE_COMPLEX * NDIRAC * NCOL * volh_tot

  allocate(x(n)); call move_alloc(x, cg_arrays%r)
  allocate(x(n)); call move_alloc(x, cg_arrays%p)
  allocate(x(n)); call move_alloc(x, cg_arrays%aap)
  allocate(x(n)); call move_alloc(x, cg_arrays%tmp1)
  allocate(x(n)); call move_alloc(x, cg_arrays%tmp2)
  allocate(x(SIZE_COMPLEX * NDIRAC2 * NCOL * volh_tot * DIM * FWD_BWD)); call move_alloc(x, cg_arrays%a)

  select case (input%tuning_cg_precision)
  case (32)
     call cg_zero_r4(n, cg_arrays%r)
     call cg_zero_r4(n, cg_arrays%p)
     call cg_zero_r4(n, cg_arrays%aap)
     call cg_zero_r4(n, cg_arrays%tmp1)
     call cg_zero_r4(n, cg_arrays%tmp2)
  case (64)
     call cg_zero_r8(n, cg_arrays%r)
     call cg_zero_r8(n, cg_arrays%p)
     call cg_zero_r8(n, cg_arrays%aap)
     call cg_zero_r8(n, cg_arrays%tmp1)
     call cg_zero_r8(n, cg_arrays%tmp2)
  case default
     call die("init_cg_arrays(): unknown tuning_cg_precision")
  end select

end

!-------------------------------------------------------------------------------
subroutine init_module_cg_conf()

  use module_cg_conf
  use module_input
  use module_vol
  implicit none

  real(8), dimension(:), allocatable :: x
  integer :: n  
  ATTRIBUTE_ALIGN(x)

  n = SIZE_COMPLEX * NDIRAC * NCOL * volh_tot

  allocate(x(SIZE_COMPLEX * NCOL * NCOL * volh_tot * EVEN_ODD * (DIM+1) * FWD_BWD)); call move_alloc(x, cg_conf%u)
  allocate(x(n)); call move_alloc(x, cg_conf%in)
  allocate(x(n)); call move_alloc(x, cg_conf%out)
  allocate(x(SIZE_COMPLEX * 18 * 2 * volh)); call move_alloc(x, cg_conf%c_even)
  allocate(x(SIZE_COMPLEX * 18 * 2 * volh)); call move_alloc(x, cg_conf%c_odd)

  select case (input%tuning_cg_precision)
  case (32)
     call cg_zero_r4(n, cg_conf%in)
     call cg_zero_r4(n, cg_conf%out)
  case (64)
     call cg_zero_r8(n, cg_conf%in)
     call cg_zero_r8(n, cg_conf%out)
  case default
     call die("init_cg_confs(): unknown tuning_cg_precision")
  end select

end

!-------------------------------------------------------------------------------
subroutine init_module_cg_proj_xbound()

  use module_cg_proj_xbound
  use module_surface
  implicit none
  integer :: mu
  integer :: blocklength_r4, blocklength_r8, blocklength_r4s, blocklength_r8s
!!  integer :: type_spinor2, type_spinor4, lb, extent, ierror
  integer :: type_spinor2, type_spinor4, ierror
  integer(kind=MPI_ADDRESS_KIND) :: lb, extent

  blocklength_r4  = SIZE_COMPLEX * NDIRAC2
  blocklength_r8  = SIZE_COMPLEX * NDIRAC2
  blocklength_r4s = SIZE_COMPLEX * NDIRAC2 * SIMD_WIDTH_R4
  blocklength_r8s = SIZE_COMPLEX * NDIRAC2 * SIMD_WIDTH_R8

  do mu = 2, DIM
     xbound_count_proj(mu)   = surface2(mu)%nsites * SIZE_COMPLEX * NDIRAC2 * NCOL

     xbound_count_sc_r4(mu)  = xbound_count_proj(mu) / blocklength_r4
     xbound_count_sc_r8(mu)  = xbound_count_proj(mu) / blocklength_r8
     xbound_count_sc_r4s(mu) = xbound_count_proj(mu) / blocklength_r4s
     xbound_count_sc_r8s(mu) = xbound_count_proj(mu) / blocklength_r8s
  enddo

  FOR rkind in 4 8
  FOR simd in '' s
  SET _rrs = _r %% rkind %% simd
  call mpi_type_contiguous(2 * blocklength%%_rrs, MPI_REAL%%rkind, type_spinor4, ierror)  
  call mpi_type_commit(type_spinor4, ierror)
  call mpi_type_get_extent(type_spinor4, lb, extent, ierror)
  call mpi_type_free(type_spinor4, ierror)
  call mpi_type_vector(1, blocklength%%_rrs, 2 * blocklength%%_rrs, MPI_REAL%%rkind, type_spinor2, ierror)
  call mpi_type_create_resized(type_spinor2, lb, extent, xbound_type_sc%%_rrs, ierror)
  call mpi_type_commit(xbound_type_sc%%_rrs, ierror)

  ENDSET
  ENDFOR
  ENDFOR

  xbound_request(:, :, :) = MPI_REQUEST_NULL
  
end

!------------------------------------------------------------------------------
subroutine init_module_cg_threads()

  use module_cg_threads
  use module_cg_vol
  use module_function_decl
  use module_thread
  implicit none
  integer :: i
 
  CHECK(num_threads() <= cg_max_threads)
  if (num_threads() < 2) return

  FOR _rrs in _r4 _r8 _r4s _r8s

  if (num_threads() == 1) then
     i_start%%_rrs(0) = 1
     i_end%%_rrs(0) = volh%%_rrs
  else
     i_start%%_rrs(0) = 0
     i_end%%_rrs(0) = -1
     call init_common_thread_split(num_threads() - 1, volh%%_rrs, i_start%%_rrs(1:), i_end%%_rrs(1:))
  endif
  ENDFOR

end

!===============================================================================
