!===============================================================================
!
! ildg_conf.F90 - read/write configurations in ILDG format
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2007-2017 Hinnerk Stueben
!               2010, 2013 Yoshifumi Nakamura
!
! 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/>.
!
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
module module_ildg_conf

  character(*), parameter :: READ = "read"
  character(*), parameter :: WRITE = "write"
  character(*), parameter :: LFN_KEY = "bqcd-restart"

  character, parameter    :: ascii_null = achar(0)

  character(80)           :: prefix
  character(80)           :: middle
  character(256)          :: lfn
  integer                 :: lx, ly, lz, lt, precision
  INT8                    :: words, bytes, cksum, su3_cksum
  INT8                    :: u1_words, u1_bytes, u1_cksum
  INT8                    :: axion_words, axion_bytes, axion_cksum
  REAL                    :: plaquette
  REAL                    :: plaquette_qed
  REAL                    :: sk_axion_ena, sk_axion_ave
end

!-------------------------------------------------------------------------------
subroutine conf_read_ildg(restart, para, conf)

  use typedef_para
  use module_function_decl
  use module_input
  use module_p_interface
  use module_vol
  use module_ildg_conf
  use module_switches
  implicit none

  logical                               :: restart
  type(type_para)                       :: para
  type(hmc_conf), dimension(MAX_TEMPER) :: conf

  P_GAUGE_FIELD_ILDG, save              :: u_ildg
  P_U1_GENERATOR_FIELD_ILDG, save       :: u1_ildg
  P_AXION_FIELD_ILDG, save              :: a_ildg
  FILENAME                              :: filename
  integer                               :: len_lfn
  INT8                                  :: bytes_read, bytes_total
  integer :: i, mu, eo
  REAL, external                        :: sg, sg_qed, sk_axion
  REAL                                  :: dummy

  call begin(UREC, "ILDGread")

  ALLOCATE_G_FIELD_ILDG(u_ildg)

  call ildg_lime_file(filename, READ, restart)
#if  defined(LEMON) && !defined(MPI_1)
  call ildg_open_r_lemon(filename, lx, ly, lz, lt, para%NPE, precision, bytes)
  call ildg_check()
  call ildg_set_sizes()
  TIMING_START(timing_bin_u_read_ildg)
  call ildg_io_lemon(READ, u_ildg, precision)
  TIMING_STOP(timing_bin_u_read_ildg)
  len_lfn = len(lfn) - 1
  call ildg_close_r_lemon(cksum, bytes_read, lfn, len_lfn)
#else
  call ildg_open_r(filename, lx, ly, lz, lt, precision, bytes)
  call ildg_check()
  call ildg_set_sizes()
  
  TIMING_START(timing_bin_u_read_ildg)
  call ildg_io(READ, u_ildg)
  TIMING_STOP(timing_bin_u_read_ildg)

  if (switches%qed) then
     ALLOCATE_U1GEN_FIELD_ILDG(u1_ildg)
     call ildg_find_u1_binary_data()
     call ildg_u1_io(READ, u1_ildg)
     call ildg_u1_seq(READ, conf(1)%a_qed, u1_ildg)
     call xbound_u1gen_field(conf(1)%a_qed)
     dummy = sg_qed(conf(1)%a_qed, plaquette_qed)
  else
     plaquette_qed = ZERO
  endif

  if (switches%axion) then
     ALLOCATE_AXION_FIELD_ILDG(a_ildg)
     call ildg_find_axion_binary_data()
     call ildg_axion_io(READ, a_ildg)
     call ildg_axion_seq(READ, conf(1)%axion, a_ildg)
     call xbound_axion_field(conf(1)%axion)
     dummy = sk_axion(conf(1)%axion, sk_axion_ena, sk_axion_ave)
  else
     sk_axion_ena = ZERO
  endif

  len_lfn = len(lfn) - 1
  call ildg_close_r(cksum, bytes_read, lfn, len_lfn)
#endif

  call ildg_seq(READ, conf(1)%u, u_ildg)

  !!!call conf_normalize(conf(1)%u)

  if (precision == 32) call conf_check32(conf(1)%u)
  if (precision == 64) call conf_check(conf(1)%u)

  call xbound_g_field(conf(1)%u)
  conf(1)%former = 1

  plaquette = ONE - (sg(conf(1)%u) / (SIX * volume))

  if (my_pe() == 0) then
     call ildg_check2(restart)
     call ildg_rec(READ, filename)

                         bytes_total = bytes
     if (switches%qed)   bytes_total = bytes_total + u1_bytes
     if (switches%axion) bytes_total = bytes_total + axion_bytes

     ASSERT(bytes_total == bytes_read)
  endif

  call end(UREC, "ILDGread")
end

!-------------------------------------------------------------------------------
subroutine conf_write_ildg(restart, para, conf)

  use typedef_para
  use module_function_decl
  use module_input
  use module_p_interface
  use module_vol
  use module_ildg_conf
  use module_switches
  implicit none

  logical                               :: restart
  type(type_para)                       :: para
  type(hmc_conf), dimension(MAX_TEMPER) :: conf

  P_GAUGE_FIELD_ILDG, save              :: u_ildg
  P_U1_GENERATOR_FIELD_ILDG, save       :: u1_ildg
  P_AXION_FIELD_ILDG, save              :: a_ildg
  FILENAME                              :: filename
  INT8                                  :: cksum0, bytes_written, bytes_total
  integer                               :: i, z, t
  REAL, external                        :: sg, sg_qed, sk_axion
  REAL                                  :: dummy

  call begin(UREC, "ILDGwrite")

  ALLOCATE_G_FIELD_ILDG(u_ildg)

  plaquette  = ONE - (sg(conf(1)%u) / (SIX * volume))
  plaquette_qed = ZERO

  call ildg_set_prefix()
  call ildg_set_middle()
  call ildg_lime_file(filename, WRITE, restart)

  lx = para%l(1)
  ly = para%l(2)
  lz = para%l(3)
  lt = para%l(4)
  
  if (restart) then
     precision = input%ildg_precision_restart
  else
     precision = input%ildg_precision
  endif

  call ildg_set_sizes()

  call ildg_check()
  call ildg_seq(WRITE, conf(1)%u, u_ildg)


#if  defined(LEMON) && !defined(MPI_1)
  call ildg_open_w_lemon(filename, lx, ly, lz, lt, para%NPE, precision, bytes)
  TIMING_START(timing_bin_u_write_ildg)
  call ildg_io_lemon(WRITE, u_ildg, precision)
  TIMING_STOP(timing_bin_u_write_ildg)
  call cksum_get(cksum, bytes_written)
  call ildg_set_lfn(filename, restart)
  call ildg_close_w_lemon(cksum0, bytes_written, lfn)

  if (my_pe() == 0) then
     call ildg_rec(WRITE, filename)
     if (.not. restart) then
        i = index(lfn, ascii_null)
        ASSERT(i > 0)
 call ildg_meta("ensemble", para, conf, lfn(1:i-1), plaquette, precision, su3_cksum)
 call ildg_meta("conf", para, conf, lfn(1:i-1), plaquette, precision, su3_cksum)
     endif
     ASSERT(cksum0 == cksum)
     ASSERT(bytes == bytes_written)
  endif

#else

  if (my_pe() == 0) then
     call ildg_open_w(filename, lx, ly, lz, lt, precision, bytes)
  endif

  TIMING_START(timing_bin_u_write_ildg)
  call ildg_io(WRITE, u_ildg)
  TIMING_STOP(timing_bin_u_write_ildg)

  if (switches%qed) then
     dummy = sg_qed(conf(1)%a_qed, plaquette_qed)
     ALLOCATE_U1GEN_FIELD_ILDG(u1_ildg)
     call ildg_u1_seq(WRITE, conf(1)%a_qed, u1_ildg)
     if (my_pe() == 0) then
        call qcdsf_u1_w(lx, ly, lz, lt, precision, u1_bytes)
     endif
     call ildg_u1_io(WRITE, u1_ildg)
  endif

  if (switches%axion) then
     ALLOCATE_AXION_FIELD_ILDG(a_ildg)
     dummy = sk_axion(conf(1)%axion, sk_axion_ena, sk_axion_ave)
     call ildg_axion_seq(WRITE, conf(1)%axion, a_ildg)
     if (my_pe() == 0) then
        call qcdsf_axion_w(lx, ly, lz, lt, precision, axion_bytes)
     endif
     call ildg_axion_io(WRITE, a_ildg)
  endif

  if (my_pe() == 0) then
     call cksum_get(cksum, bytes_written)
     call ildg_set_lfn(filename, restart)
  endif

  if (my_pe() == 0) then
     call ildg_close_w(cksum0, bytes_written, lfn)
     ! ildg_close_w() calls ranfield_dump_ildg() as well 
  else
     do t = 0, lt - 1
     do z = 0, lz - 1
        call ranfield_dump_ildg(i, 0, z, t)  ! "i" is dummy
     enddo
     enddo
  endif

  if (my_pe() == 0) then
     call ildg_rec(WRITE, filename)

     if (.not. restart) then
        i = index(lfn, ascii_null)
        ASSERT(i > 0)
 call ildg_meta("ensemble", para, conf, lfn(1:i-1), plaquette, precision, su3_cksum)
 call ildg_meta("conf", para, conf, lfn(1:i-1), plaquette, precision, su3_cksum)
     endif

                         bytes_total = bytes
     if (switches%qed)   bytes_total = bytes_total + u1_bytes
     if (switches%axion) bytes_total = bytes_total + axion_bytes

     ASSERT(cksum0 == cksum)
     ASSERT(bytes_total == bytes_written)
  endif

#endif
  call end(UREC, "ILDGwrite")
end

!-------------------------------------------------------------------------------
subroutine ildg_set_sizes()

  use module_ildg_conf
  use module_vol
  implicit none
  integer :: bytes_per_word

  select case (precision)
     case(32)
        bytes_per_word = 4
     case(64)
        bytes_per_word = 8
     case default
        call die("ildg_set_sizes(): unknown precision")
  end select

  words = SIZE_COMPLEX * NCOL * NCOL * DIM 
  words = words * volume                   ! 64 bit integer arithmetic
  bytes = words * bytes_per_word

  u1_words = DIM
  u1_words = u1_words * volume
  u1_bytes = u1_words * bytes_per_word

  axion_words = 1
  axion_words = axion_words * volume
  axion_bytes = axion_words * bytes_per_word

  su3_cksum = 0
  u1_cksum = 0
  axion_cksum = 0

end  

!-------------------------------------------------------------------------------
subroutine ildg_check()

  use module_ildg_conf
  use module_input
  implicit none
  INT8 :: bytes0

  CHECK(input%lattice(1) == lx)
  CHECK(input%lattice(2) == ly)
  CHECK(input%lattice(3) == lz)
  CHECK(input%lattice(4) == lt)
  CHECK(precision == 32 .or. precision == 64)

  bytes0 = bytes
  call ildg_set_sizes()
  CHECK(bytes == bytes0)

end

!-------------------------------------------------------------------------------
subroutine ildg_check2(restart)

  use module_counter
  use module_ildg_conf
  implicit none
  logical, intent(in) :: restart

  character :: a, b
  integer   :: run, job, traj, i
  INT8      :: cksum0, bytes0

  if (index(lfn, LFN_KEY) > 0) then
     i = index(lfn, ascii_null)
     ASSERT(i > 0)
     read(lfn(1:i-1), *) a, b, cksum0, bytes0, run, job, traj
     CHECK(cksum == cksum0)
     if (restart) then
        CHECK(run == counter%run)
        CHECK(job + 1 == counter%job)
        CHECK(traj == counter%traj)
     endif
  endif

end

!-------------------------------------------------------------------------------
subroutine ildg_filename(fname, action, restart, extension)

  use module_bqcd
  use module_counter
  use module_ildg_conf
  use module_input
  implicit none
  character(*), intent(out) :: fname
  character(*), intent(in)  :: action
  logical, intent(in)       :: restart
  character(*), intent(in)  :: extension
  character(32)             :: fmt_restart
  integer                   :: i

  if (restart) then
     if (input%run < 1000) then
        fmt_restart = "(a, '.', i3.3, '.', 2a)"
     else
        fmt_restart = "(a, '.', i0, '.', 2a)"
     endif
     write(fname, fmt_restart) &
           trim(prog_name), input%run, trim(extension), ascii_null
  else
     if (action == READ) then
        if (input%start_ildg_file == "") then
           call die("ildg_filename(): start_ildg_file not set")
        else
           fname = input%start_ildg_file
           i = index(fname, " ")
           CHECK(i > 0)
           CHECK(i <= len(fname))
           fname(i:i) = ascii_null
        endif
     elseif (action == WRITE) then
        write(fname, "(5a)") &
              trim(prefix), trim(middle), ".", trim(extension), ascii_null
     else
        call die("ildg_filename(): unknown action")
     endif
  endif
end

!-------------------------------------------------------------------------------
subroutine ildg_lime_file(fname, action, restart)

  use module_input
  implicit none
  character(*), intent(out) :: fname
  character(*), intent(in)  :: action
  logical, intent(in)       :: restart

  call ildg_filename(fname, action, restart, input%ildg_filename_extension)
end

!-------------------------------------------------------------------------------
subroutine ildg_set_lfn(fname, restart)

  use module_counter
  use module_input
  use module_ildg_conf
  implicit none

  character(*), intent(in)  :: fname
  logical, intent(in)       :: restart
  integer                   :: i

  if (restart) then

     ! save counters and checksum in the lime file instead of lfn

     i = index(fname, ascii_null)

     write(lfn, *) LFN_KEY, " ", &
                   fname(1:i-1), &
                   cksum, bytes, &
                   counter%run, counter%job, counter%traj, &
                   ascii_null
  else
     call ildg_expand_input_lfn_mc()
     lfn = trim(input%ildg_data_lfn_path) // "/" // trim(fname) // ascii_null
  endif

end

!-------------------------------------------------------------------------------
subroutine ildg_seq(action, u, u_out)

  use module_ildg_conf
  use module_lattice
  use module_vol
  implicit none
  character(*), intent(in)      :: action
  GAUGE_FIELD, intent(in)       :: u
  GAUGE_FIELD_ILDG, intent(out) :: u_out

  if (precision == 32) call ildg_seq_r4(action, u, u_out)
  if (precision == 64) call ildg_seq_r8(action, u, u_out)

end

!-------------------------------------------------------------------------------
subroutine ildg_u1_seq(action, u, u_out)

  use module_ildg_conf
  use module_lattice
  use module_vol
  implicit none
  character(*), intent(in)      :: action
  U1_GENERATOR_FIELD, intent(in)       :: u
  U1_GENERATOR_FIELD_ILDG, intent(out) :: u_out

  if (precision == 32) call ildg_u1_seq_r4(action, u, u_out)
  if (precision == 64) call ildg_u1_seq_r8(action, u, u_out)

end

!-------------------------------------------------------------------------------
subroutine ildg_axion_seq(action, a, a_out)

  use module_ildg_conf
  use module_lattice
  use module_vol
  implicit none
  character(*), intent(in)      :: action
  AXION_FIELD, intent(in)       :: a
  AXION_FIELD_ILDG, intent(out) :: a_out

  if (precision == 32) call ildg_axion_seq_r4(action, a, a_out)
  if (precision == 64) call ildg_axion_seq_r8(action, a, a_out)

end

!-------------------------------------------------------------------------------
subroutine ildg_io(action, u)

  use module_ildg_conf
  use module_lattice_io
  use module_vol
  implicit none
  character(*), intent(in)        :: action
  GAUGE_FIELD_ILDG, intent(inout) :: u

  integer, parameter :: words_per_site = SIZE_COMPLEX * NCOL * NCOL * DIM

  if (precision == 32) call ildg_io_r4(action, u, words_per_site, su3_cksum)
  if (precision == 64) call ildg_io_r8(action, u, words_per_site, su3_cksum)
end

!-------------------------------------------------------------------------------
subroutine ildg_u1_io(action, u)

  use module_ildg_conf
  use module_lattice_io
  use module_vol
  implicit none
  character(*), intent(in)               :: action
  U1_GENERATOR_FIELD_ILDG, intent(inout) :: u

  integer, parameter :: words_per_site = DIM

  if (precision == 32) call ildg_io_r4(action, u, words_per_site, u1_cksum)
  if (precision == 64) call ildg_io_r8(action, u, words_per_site, u1_cksum)
end

!-------------------------------------------------------------------------------
subroutine ildg_axion_io(action, a)

  use module_ildg_conf
  use module_lattice_io
  use module_vol
  implicit none
  character(*), intent(in)        :: action
  AXION_FIELD_ILDG, intent(inout) :: a

  integer, parameter :: words_per_site = 1

  if (precision == 32) call ildg_io_r4(action, a, words_per_site, axion_cksum)
  if (precision == 64) call ildg_io_r8(action, a, words_per_site, axion_cksum)
end

!-------------------------------------------------------------------------------
subroutine ildg_rec(action, fname)

  use module_ildg_conf
  use module_switches
  implicit none
  character(*), intent(in) :: action
  character(*), intent(in) :: fname

  character(*), parameter :: fmt = "(x,4a)"
  integer :: i, j
  INT8 :: bytes_total

                      bytes_total = bytes
  if (switches%qed)   bytes_total = bytes_total + u1_bytes
  if (switches%axion) bytes_total = bytes_total + axion_bytes

  i = index(fname, ascii_null); ASSERT(i > 0)
  j = index(lfn, ascii_null);   ASSERT(j > 0)

                      write(UREC, fmt) "ildg-", action, " file ", fname(1:i-1)
                      write(UREC, *)   "ildg-", action, " precision ", precision
                      write(UREC, *)   "ildg-", action, " plaquette     ", plaquette
  if (switches%qed)   write(UREC, *)   "ildg-", action, " plaquette_qed ", plaquette_qed
  if (switches%axion) write(UREC, *)   "ildg-", action, " sk_axion_ena  ", sk_axion_ena

                      write(UREC, *)   "ildg-", action, " bytes       ", bytes
  if (switches%qed)   write(UREC, *)   "ildg-", action, " bytes_u1    ", u1_bytes
  if (switches%axion) write(UREC, *)   "ildg-", action, " bytes_axion ", axion_bytes
                      write(UREC, *)   "ildg-", action, " bytes_total ", bytes_total

                      write(UREC, *)   "ildg-", action, " cksum       ", su3_cksum
  if (switches%qed)   write(UREC, *)   "ildg-", action, " cksum_u1    ", u1_cksum
  if (switches%axion) write(UREC, *)   "ildg-", action, " cksum_axion ", axion_cksum
                      write(UREC, *)   "ildg-", action, " cksum_total ", cksum

                      write(UREC, fmt) "ildg-", action, " lfn ", lfn(1:j-1)
end

!-------------------------------------------------------------------------------
subroutine ildg_set_prefix

  use module_ildg_conf
  use module_input
  implicit none
  
  input%ildg_filename_prefix = adjustl(input%ildg_filename_prefix)

  call expand_filename(input%ildg_filename_prefix, prefix)
end

!-------------------------------------------------------------------------------
subroutine ildg_set_middle

  use module_ildg_conf
  use module_input
  implicit none
  
  input%ildg_filename_middle = adjustl(input%ildg_filename_middle)

  call expand_filename(input%ildg_filename_middle, middle)
end

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