!===============================================================================
!
! cg_mre_matrix.F90
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2006-2010 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/>.
!
!-------------------------------------------------------------------------------
! Added MRE for custom matrices, which is used in the solver_cg_poly=1 case.
!
! Editor: Taylor Haar
! Last modified: 29/09/2016
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
module typedef_mre_matrix
  use module_mre
  type(type_mre), dimension(10), save :: solutions
end

!-------------------------------------------------------------------------------
subroutine cg_mre_matrix(a, x, id, np, calc_sf, matrix, half_matrix)
  ! MRE for a custom matrix: assumes matrix = half_mat^{dag} half_mat

  use typedef_mre_matrix
  use module_vol
  implicit none

  integer,        intent(in)  :: calc_sf, id, np ! np = n-th pseudo fermion
  SPINCOL_FIELD,  intent(in)  :: x
  SPINCOL_FIELD,  intent(out) :: a

  external :: matrix, half_matrix

  DEBUG2S("Start: cg_mre_matrix")

  call mre_get2_matrix(solutions(np), a, x, id, np, half_matrix)
  call cg_outer_matrix(a, x, id, np, calc_sf, matrix)
  call mre_put(solutions(np), a, calc_sf)

  DEBUG2S("End: cg_mre_matrix")
end

!-------------------------------------------------------------------------------
subroutine mre_get2_matrix(basis, trial, phi, id, mid, half_matrix)

  ! get trial solution for custom matrix = half_mat^{dag} half_mat

  use module_function_decl
  use module_mre
  use module_vol
  implicit none

  type(type_mre), intent(inout) :: basis
  SPINCOL_FIELD,  intent(out)   :: trial
  SPINCOL_FIELD,  intent(in)    :: phi
  integer , intent(in) :: id, mid

  external :: half_matrix

  type(type_mre), save          :: mv  ! "Matrix * v"
  COMPLEX, target               :: g(mre_n_vec, mre_n_vec + 1)
  COMPLEX, pointer              :: b(:)
  integer                       :: size_g
  integer                       :: i
  integer                       :: j
  integer                       :: s
  integer                       :: c
  integer                       :: rest


  DEBUG2S("Start: mre_get2")

  if (mre_n_vec == 0 .or. .not. associated(basis%vec) .or. basis%rank == 0) then
     call sc_copy(trial, phi)
     DEBUG2S("End: mre_get2")
     return
  endif

  if (basis%rank == 1) then
     call sc_copy(trial, basis%vec(1)%sc(1,1,1))
     DEBUG2S("End: mre_get2")
     return
  endif

  size_g = mre_n_vec * (mre_n_vec + 1) * SIZE_COMPLEX

  b => g(:, mre_n_vec + 1)  ! storage arrangement for global sum
                            ! => one global sum for everything

  call mre_allocate(mv)
  call mre_gram_schmidt(basis)

  do i = 1, basis%rank
     b(i) = sc_cdotc(basis%vec(i)%sc(1,1,1), phi)
     call half_matrix(mv%vec(i)%sc(1,1,1), basis%vec(i)%sc(1,1,1), id, mid)
  enddo

  do i = 1, basis%rank
     g(i, i) = sc_norm2(mv%vec(i)%sc(1,1,1))
     do j = i + 1, basis%rank
        g(i, j) = sc_cdotc(mv%vec(i)%sc(1,1,1), mv%vec(j)%sc(1,1,1))
        g(j, i) = conjg(g(i, j))
     enddo
  enddo

  call global_sum_vec(size_g, g)

  call mre_gauss_jordan(g, b, basis%rank, mre_n_vec)

  ! calculation of "trial" with doubled data re-use:

  call sc_cax2(trial, basis%vec(1)%sc(1,1,1), b(1), basis%vec(2)%sc(1,1,1), b(2))

  rest = mod(basis%rank, 2)

  do j = 3, basis%rank - rest, 2
     call sc_caxpy2(trial, basis%vec(j)%sc(1,1,1),   b(j), &
                           basis%vec(j+1)%sc(1,1,1), b(j+1))
  enddo

  if (rest == 1) then
     j = basis%rank
     call sc_caxpy(trial, basis%vec(j)%sc(1,1,1), b(j))
  endif

  DEBUG2S("End: mre_get2")

end

