module big_block_mod
  use elemental_mod
  use geometry

implicit none

  !> big matrix block, to each element \f$ K\in{\cal T}_h \f$ corresponds several blocks
  type, public :: BigBlock_t
     real, allocatable, dimension (:,:,:,:,:,:)  :: Mb  ! block of matrix in R^6

     contains

     !procedure :: copyBigBlock
     procedure :: printme
     procedure :: dimensions
     procedure :: zero
     !procedure :: copyMatrixToBigBlock
     procedure :: copyMBlockToBigBlock
     procedure :: copyBigBlockToMBlock

!     procedure :: delete

  end type BigBlock_t

  interface BigBlock_t
       procedure :: constructSquareFromDimensions
       procedure :: constructRectangleFromDimensions
       procedure :: constructFromMatrix
  end interface

  public:: operator(+), operator(-), operator(*)

  interface operator(+)
    module procedure addBigBlock
  end interface

  interface operator(-)
    module procedure subBigBlock
  end interface

  interface operator(*)
    module procedure mulBigBlockConstant
    !module procedure mulBigBlock
  end interface

  interface assignment (=)
      procedure :: copyBigBlock
!      procedure :: copyMatrixToBigBlock
      !procedure :: copyBigBlockToMatrix

  end interface

  public :: assignment(=)

  public :: vectorMulBigBlock
  public :: vectorMulBigBlockTransposed
  public :: SolveLocalBigMatrixProblem

contains


   function dimensions( this)
      class( BigBlock_t ), intent(in) :: this
      integer, dimension(1:6) :: dimensions
      if (allocated(this%Mb)) then
      	dimensions =  (/ size(this%Mb,1), size(this%Mb,2), size(this%Mb,3), &
                       size(this%Mb,4), size(this%Mb,5), size(this%Mb,6) /)
      else
        print*, 'bigBlock matrix Mb is not allocated!'
        dimensions = (/-1, -1, -1, -1, -1, -1/)
      endif

   end function dimensions


   function constructSquareFromDimensions(dof,nd,tdof) result(this)
        type( BigBlock_t ) :: this
        integer, intent(in) :: dof, nd, tdof

        if (allocated(this%Mb)) &
            deallocate(this%Mb)

        allocate(this%Mb(dof, dof, nd,nd, tdof, tdof), source = 0.0)

   end function constructSquareFromDimensions

   function constructRectangleFromDimensions(dof1, dof2, nd1, nd2, tdof1, tdof2) result(this)
        type( BigBlock_t ) :: this
        integer, intent(in) :: dof1, dof2, nd1, nd2, tdof1, tdof2

        if (allocated(this%Mb)) &
            deallocate(this%Mb)

        allocate(this%Mb(dof1,dof2, nd1,nd2, tdof1, tdof2), source = 0.0)

   end function constructRectangleFromDimensions

   function constructFromMatrix(matrix) result(this)
        type( BigBlock_t ) :: this
        real, dimension(:,:,:,:,:,:), intent(in) :: matrix
        integer :: dof1, dof2, nd1, nd2, tdof1, tdof2

        dof1 = size(matrix ,1)
        dof2 = size(matrix ,2)

        nd1  = size(matrix ,3)
        nd2  = size(matrix ,4)

        tdof1  = size(matrix ,5)
        tdof2  = size(matrix ,6)

        if (allocated(this%Mb)) &
            deallocate(this%Mb)

        allocate(this%Mb(dof1,dof2,nd1,nd2,tdof1,tdof2), source = matrix)

   end function constructFromMatrix

   subroutine printme(this)
      class(BigBlock_t), intent(inout) :: this

      write(*,*) 'Not done'

   end subroutine printme


   subroutine zero(this)
      class(BigBlock_t), intent(inout) :: this

      this%Mb(:,:,:,:,:,:) = 0.0

   end subroutine zero


   subroutine copyBigBlock(this,other)
      class(BigBlock_t),intent(inout) :: this
      class(BigBlock_t),intent(in) :: other

      if (allocated(other%Mb)) then
        this%Mb = other%Mb
      else
        stop 'copy BigBlock from unallocated matrix!'
      end if
   end subroutine copyBigBlock

!   subroutine copyMatrixToBigBlock(this,matrix)
!      class(BigBlock_t),intent(inout) :: this
!      real, dimension(:,:,:,:,:,:), intent(in) :: matrix
!
!      !this%constructFromMatrix(matrix)
!      this%Mb = matrix
!
!   end subroutine copyMatrixToBigBlock
!   ! we suppose that the matrix has the structure used eg for blockST -
!   ! 1:dof,1:ndim,1:Tdof
!   subroutine copyMatrixToBigBlock(this,matrix, dimensions)
!      class(BigBlock_t),intent(inout) :: this
!      real, dimension(:,:), intent(in) :: matrix
!      integer, dimension(1:6),intent(in) :: dimensions
!      integer :: row1, row2, row3, column1, column2, column3
!      integer :: i,j,m,n
!      integer :: longRow, longColumn
!
!      !this%constructFromMatrix(matrix)
!      !this%Mb = matrix
!      row1 = dimensions(1)
!      column1 = dimensions(2)
!
!      row2 = dimensions(3)
!      column2 = dimensions(4)
!
!      row3 = dimensions(5)
!      column3 = dimensions(6)
!
!      print*, 'control what does the if ANY do?'
!      if ( any( this%dimensions() /= dimensions ) ) &
!        stop 'wrong dimensions in copySquareMatrixToBigBlock'
!
!
!      longRow = 1
!      longColumn = 1
!
!      do i = 1,row3
!        do m = 1, row2
!          do j = 1, column3
!            do n = 1, column2
!                this%Mb(1:row1,1:column1, i, j, m, n) = &
!                  matrix( longRow: longRow + row1 - 1, longColumn:longColumn + column1 - 1 )
!               longColumn = longColumn + column1
!            end do ! n
!          end do ! j
!          longRow = longRow + row1
!          longColumn = 1
!        end do ! m
!      end do ! i
!
!
!   end subroutine copyMatrixToBigBlock

   ! we suppose that the matrix has the structure used eg for blockST -
   ! 1:dof,1:ndim,1:Tdof
   subroutine copyMBlockToBigBlock(this, matrix, dimensions)
      class(BigBlock_t),intent(inout) :: this
      type( Mblock) , intent(in) :: matrix
      !real, dimension(:,:), intent(in) :: matrix
      integer, dimension(1:6),intent(in) :: dimensions
      integer :: row1, row2, row3, column1, column2, column3
      integer :: i,j,m,n
      integer :: longRow, longColumn

      !this%constructFromMatrix(matrix)
      !this%Mb = matrix
      row1 = dimensions(1)
      column1 = dimensions(2)

      row2 = dimensions(3)
      column2 = dimensions(4)

      row3 = dimensions(5)
      column3 = dimensions(6)


      !print*, 'control what does the if ANY do?' , any( this%dimensions() < dimensions )
      if ( any( this%dimensions() < dimensions ) ) then
        print*, 'dimensions of block:', dimensions
        print*, 'dimensions of bigBlock:', this%dimensions()
        stop 'wrong dimensions in copySquareMatrixToBigBlock'
      endif

      longRow = 1
      longColumn = 1

      do i = 1,row3
        do m = 1, row2
          do j = 1, column3
            do n = 1, column2

                this%Mb(1:row1,1:column1, m, n, i, j) = &
                  matrix%Mb( longRow: longRow + row1 - 1, longColumn:longColumn + column1 - 1 )
               longColumn = longColumn + column1
            end do ! n
          end do ! j
          longRow = longRow + row1
          longColumn = 1
        end do ! m
      end do ! i


   end subroutine copyMBlockToBigBlock

     ! we suppose that the matrix has the structure used eg for blockST -
   ! 1:dof,1:ndim,1:Tdof
   subroutine copyBigBlockToMBlock(this, matrix, dimensions)
      class(BigBlock_t),intent(in) :: this
      type( Mblock) , intent(inout) :: matrix
      integer, dimension(1:6),intent(in) :: dimensions
      integer :: row1, row2, row3, column1, column2, column3
      integer :: i,j,m,n
      integer :: longRow, longColumn

      row1 = dimensions(1)
      column1 = dimensions(2)

      row2 = dimensions(3)
      column2 = dimensions(4)

      row3 = dimensions(5)
      column3 = dimensions(6)

      if (.not. allocated(matrix%Mb) ) &
        stop 'matrix%Mb not allocated in copyBigBlockToMBlock'
      if (.not. allocated(this%Mb) ) &
        stop 'this%Mb not allocated in copyBigBlockToMBlock'

      if ( size(matrix%Mb,1) < row1*row2*row3 .or. size(matrix%Mb,2) < column1*column2*column3 )  then
        print*, 'Size of Mblock = ', size(matrix%Mb,1) , ' x ', size(matrix%Mb,2)
        print*, 'Input dimensions = ' , dimensions
        stop "matrix%Mb is too small in copyBigBlockToMBlock"
      endif

      if ( any( this%dimensions() < dimensions ) ) then
        print*, 'dimensions of block:', dimensions
        print*, 'dimensions of bigBlock:', this%dimensions()
        stop 'wrong dimensions in copyBigBlockToMBlock'
      endif

      matrix%Mb(:,:) = 0.0

      longRow = 1
      longColumn = 1

      do i = 1,row3
        do m = 1, row2
          do j = 1, column3
            do n = 1, column2
              matrix%Mb( longRow: longRow + row1 - 1, longColumn:longColumn + column1 - 1 ) = &
                this%Mb(1:row1,1:column1, m, n, i, j)
              longColumn = longColumn + column1
            end do ! n
          end do ! j
          longRow = longRow + row1
          longColumn = 1
        end do ! m
      end do ! i

   end subroutine copyBigBlockToMBlock

   subroutine deleteBigBlock(this)
      class( BigBlock_t ) :: this
      if (allocated(this%Mb)) &
          deallocate(this%Mb)

   end subroutine deleteBigBlock


  function addBigBlock(a,b) result(c)
    type(BigBlock_t),intent(in):: a,b
    type(BigBlock_t):: c

    ! control sizes ?
    !if (allocated(c%Mb)) deallocate(c%Mb)

    !allocate(c%Mb, source = a)
    c%Mb = a%Mb + b%Mb

  end function addBigBlock

  function subBigBlock(a,b) result(c)
    type(BigBlock_t),intent(in):: a,b
    type(BigBlock_t):: c

    ! control sizes ?
    !if (allocated(c%Mb)) deallocate(c%Mb)

    !allocate(c%Mb, source = a)
    c%Mb = a%Mb - b%Mb

  end function subBigBlock

  function mulBigBlockConstant(a,b) result(c)
    real,intent(in):: a
    type(BigBlock_t),intent(in):: b
    type(BigBlock_t):: c

    ! control sizes ?
    !if (allocated(c%Mb)) deallocate(c%Mb)

    !allocate(c%Mb, source = b%Mb)
    c%Mb = a * b%Mb

  end function mulBigBlockConstant

!  ! matrix multiplication
!  function mulBigBlock(a,b) result(c)
!    type(BigBlock_t),intent(in):: a,b
!    type(BigBlock_t):: c
!    integer :: Adof1, Adof2, And1, And2, Atdof1, Atdof2
!    integer :: Bdof1, Bdof2, Bnd1, Bnd2, Btdof1, Btdof2
!    integer :: dof1, dof2, nd1, nd2, tdof1, tdof2
!
!    ! control sizes
!    Adof1 = size(a ,1)
!    Adof2 = size(a ,2)
!    And1  = size(a ,3)
!    And2  = size(a ,4)
!    Atdof1  = size(a ,5)
!    Atdof2  = size(a ,6)
!
!    Bdof1 = size(b ,1)
!    Bdof2 = size(b ,2)
!    Bnd1 = size(b ,3)
!    Bnd2 = size(b ,4)
!    Btdof1  = size(b ,5)
!    Btdof2  = size(b ,6)
!
!    if (Adof2 == Bdof1 .and. And2 == Bnd1 .and. Atdof2 == Btdof1) then
!    else
!      stop 'wrong dimensions in matrix multiplication in mulBigBlock'
!    endif
!
!    allocate(c%Mb, source = b)
!    c%Mb = a * b%Mb
!
!  end function mulBigBlock

  ! matrix times vector MATMUL multiplication
  subroutine vectorMulBigBlock(A,x, dimensions, b)
    type(BigBlock_t),intent(in):: A
    type(Elemental3_t),intent(in):: x
    integer, dimension(1:6) :: dimensions
    type(Elemental3_t), intent(out) :: b
    integer :: dof1, dof2, nd1, nd2, tdof1, tdof2
    integer :: i,j , m,n
    integer, dimension(1:6) :: aDimensions

    ! controls
    aDimensions = A%dimensions()

    dof1 = dimensions(1)
    dof2 = dimensions(2)
    nd1  = dimensions(3)
    nd2  = dimensions(4)
    tdof1  = dimensions(5)
    tdof2  = dimensions(6)

    if ( any(aDimensions < dimensions) .or. &
         size(x%x,1) /= dof2 .or. size(x%x,2) /= nd2 .or. size(x%x,3) /= tdof2 ) then
       
       print*, 'dof is different = ' , dof2, size(x%x,1)
       stop 'wrong dimensions in matrix multiplication in mulBigBlock'
    end if
    
    
    call b%init( dof1, nd1, tdof1)
    
    do i = 1,tdof1
       do j = 1, nd1
          do m = 1, tdof2
             do n = 1, nd2
                !            print*, 'A = ' , A%Mb(1:dof1,1:dof2,j,n,i,m)
                !            print*, 'x = ' , x%x(1:dof2, n, m )
                !            print*,
                b%x(1:dof1,j,i) = b%x(1:dof1,j,i) + &
                     matmul( A%Mb(1:dof1,1:dof2,j,n,i,m), x%x(1:dof2, n, m ) )
             end do
          end do !m
          
       end do ! j
    end do !i
    
    
  end subroutine vectorMulBigBlock

  ! matrix TRANSPOSED times vector MATMUL multiplication
  ! FR - needs to be tested
  subroutine vectorMulBigBlockTransposed(A,x, dimensions, b)
    type(BigBlock_t),intent(in):: A
    type(Elemental3_t),intent(in):: x
    integer, dimension(1:6) :: dimensions
    type(Elemental3_t), intent(out) :: b
    integer :: dof1, dof2, nd1, nd2, tdof1, tdof2
    integer :: i,j , m,n
    integer, dimension(1:6) :: aDimensions

    ! controls
    aDimensions = A%dimensions()

    dof1 = dimensions(1)
    dof2 = dimensions(2)
    nd1  = dimensions(3)
    nd2  = dimensions(4)
    tdof1  = dimensions(5)
    tdof2  = dimensions(6)


    if ( any(aDimensions < dimensions) .or. &
      size(x%x,1) /= dof1 .or. size(x%x,2) /= nd1 .or. size(x%x,3) /= tdof1 ) then
      print*, "cscs", dimensions(7)
      print*, 'dof is different in  = vectorMulBigBlockTransposed' , dof1, size(x%x,1)
      stop 'wrong dimensions in matrix multiplication in mulBigBlock'
    end if


    call b%init( dof2, nd2, tdof2)

    do i = 1,tdof2
      do j = 1, nd2
        do m = 1, tdof1
          do n = 1, nd1
!            print*, 'A = ' , A%Mb(1:dof1,1:dof2,j,n,i,m)
!            print*, 'x = ' , x%x(1:dof2, n, m )
!            print*,
            b%x(1:dof2,j,i) = b%x(1:dof2,j,i) + &
              matmul( x%x(1:dof1, n, m ), A%Mb(1:dof1,1:dof2, n, j, m, i) )
          end do
        end do !m

      end do ! j
    end do !i


  end subroutine vectorMulBigBlockTransposed



  !> solve the problem  (\f$ A x = b \f$ ) to \f$ A\in R^{n\times n} \f$,
  !> or (\f$ A^T x = b \f$ )
  !> \f$ b\f$ represents m RHS, solution overwrites b
  subroutine SolveLocalBigMatrixProblem(n, dimensions, A, m, b, transposed)
    integer, intent(in) :: n, m
    integer, dimension(1:6) :: dimensions ! dimensions of the problem, should be < A%dimensions()
    type( BigBlock_t ), intent(in) :: A
    real, dimension(1:n, 1:m), intent(inout) :: b
    logical, intent(in) :: transposed ! A or A^T

    external:: DGESV          ! subroutines from LAPACK
    real, dimension(:), allocatable :: ipiv
    type(Mblock) :: AA
    real :: val
    integer :: k, i, info

    if (n /= dimensions(1)*dimensions(3)*dimensions(5) .or. &
      n /= dimensions(2)*dimensions(4)*dimensions(6) ) then
        stop 'Wrongly set variables in SolveLocalBigMatrixProblem'
    endif

!    if (any( A%dimensions() < dimensions )) then
!      stop 'bigBlock(0) is not big enough in RitzReconstruction'
!    endif

    allocate(ipiv(1:n))
    do k=1, n
       ipiv(k) = k
    enddo

    ! A%Mb -> Matrix AA
    !allocate( AA(1:n, 1:n)  )
    call InitMblock(AA , n , n )
    call A%copyBigBlockToMBlock(AA, dimensions)

    if ( transposed ) then
      call DGESV(n, m, transpose( AA%Mb(1:n, 1:n) ), n, ipiv, b(1:n, 1:m), n, info)
    else
      call DGESV(n, m, AA%Mb(1:n, 1:n), n, ipiv, b(1:n, 1:m), n, info)
    end if

    !call DGESV(n, m, AA%Mb(1:n, 1:n), n, ipiv, b(1:n, 1:m), n, info)
    if(INFO /= 0) print*,'Warning in lap_sub.f90 LAPACK: DGESV,  INFO = ',info

    deallocate(ipiv)
    call DeleteMblock(AA)

  end subroutine SolveLocalBigMatrixProblem


end module big_block_mod
