module solution_mod
   use mesh_mod
   use nonlinear_mod
   use define_state
   use stdgm_mod


implicit none


   public :: allocateZST
   public :: allocateZSTplus
   public :: allocateWSTplus

   public :: copyWST_fromLongVector
   public :: copyWST_toLongVector
   public :: copyZST_fromLongVector
   public :: copyZST_toLongVector
   public :: copyZST_toLongVectorMod
   public :: copyZSTplus_toLongVectorMod
   public :: copyWST_toLongVectorMod
   public :: copyWSTplus_toLongVectorMod ! not used, not tested
   public :: copyZSTplus_fromLongVector
   !public :: copyWSTplus_fromLongVector

!   public :: copyWSTtoZST
   !public :: moveZSTtoZSTplus
   public :: transfer_w_to_WS_elem
   public :: solution_L8norm

contains

   !> allocate zST and fill it with zeros
   !> if it is already allocated and the size is right, only fill with ZEROS
   ! FR_DEGPLUS - not to deg plus
   subroutine allocateZST( grid , plus)
      class( mesh ), intent(inout) :: grid
      logical, intent(in) :: plus
      class(element), pointer :: elem
      integer :: i, deg, dof, Tdof
      logical :: flag

      flag = .false.

      do i = 1, grid%nelem
         elem => grid%elem(i)
         if (plus) then
            deg = elem%deg + state%p_mod_max
            dof = DOFtriang( deg )
         else
            dof = elem%dof
         endif
         Tdof = elem%Tdof

         if ( associated(elem%zST) ) then
            if ( size(elem%zST(:,1,1)) /= ndim .or. &
                 size(elem%zST(1,:,1)) /= dof .or. &
                 size(elem%zST(1,1,:)) /= Tdof ) then
               flag = .true.
               deallocate(elem%zST)
               allocate( elem%zST(1:ndim,1:dof,1:Tdof), source = 0.0 )
            else
               elem%zST(1:ndim,1:dof,1:Tdof) = 0.0
            end if
         else
            allocate( elem%zST(1:ndim,1:dof,1:Tdof), source = 0.0 )
         endif
       end do ! i

       if (flag) then
         print*, 'allocateZST: associated elem%zST was deallocated!'
         print*
       endif

   end subroutine allocateZST

   !> allocate zST and fill it with zeros
   !> if it is already allocated and the size is right, only fill with ZEROS
   subroutine allocateZSTplus( grid )
      class( mesh ), intent(inout) :: grid
      class(element), pointer :: elem
      integer :: i, deg, dof, Tdof
      logical :: flag

      flag = .false.

      do i = 1, grid%nelem
         elem => grid%elem(i)
         deg = elem%deg + state%p_mod_max
         dof = DOFtriang( deg )
         Tdof = elem%Tdof + state%q_mod_max

         if ( associated(elem%zSTplus) ) then
            if ( size(elem%zSTplus(:,1,1)) /= ndim .or. &
                 size(elem%zSTplus(1,:,1)) /= dof ) then
               flag = .true.
               deallocate(elem%zSTplus)
               allocate( elem%zSTplus(1:ndim,1:dof,1:Tdof), source = 0.0 )
            else
               elem%zSTplus(1:ndim,1:dof,1:Tdof) = 0.0
            end if
         else
            allocate( elem%zSTplus(1:ndim,1:dof,1:Tdof), source = 0.0 )
         endif
       end do ! i

       if (flag) then
         print*, 'allocateZSTplus: associated elem%zSTplus was deallocated!'
         print*
       endif

   end subroutine allocateZSTplus

      !> allocate zST and fill it with zeros
   !> if it is already allocated and the size is right, only fill with ZEROS
   subroutine allocateWSTplus( grid )
      class( mesh ), intent(inout) :: grid
      class(element), pointer :: elem
      integer :: i, deg, dof, Tdof
      logical :: flag

      flag = .false.

      do i = 1, grid%nelem
         elem => grid%elem(i)
         deg = elem%deg + state%p_mod_max
         dof = DOFtriang( deg )
         Tdof = elem%Tdof + state%q_mod_max

         if ( associated(elem%wSTplus) ) then
            if ( size(elem%wSTplus(:,1,1)) /= ndim .or. &
                 size(elem%wSTplus(1,:,1)) /= dof ) then
               flag = .true.
               deallocate(elem%wSTplus)
               allocate( elem%wSTplus(1:ndim,1:dof,1:Tdof), source = 0.0 )
            else
               elem%wSTplus(1:ndim,1:dof,1:Tdof) = 0.0
            end if
         else
            allocate( elem%wSTplus(1:ndim,1:dof,1:Tdof), source = 0.0 )
         endif
       end do ! i

       if (flag) then
         print*, 'allocateWSTplus: associated elem%wSTplus was deallocated!'
         print*
       endif

   end subroutine allocateWSTplus

   !> distributes  elem%zST to the global vector x to
   subroutine copyZST_toLongVector( grid,  nsize, x )
      class( mesh ) , intent(in) :: grid
      integer, intent(in) :: nsize
      real, dimension(1:nsize), intent(out) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof

      kk = 0

      print*, "copyZST_toLongVector should not be called based on the old structure elem%ncv"
      print*

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (elem%ncv /= kk + 1) then
          print*, 'Problem while copying update from zST (solution.f90)'
          stop
        endif
        elemDof = elem%dof * ndim * elem%Tdof

        do k = 1, elem%Tdof
          do j = 1, ndim
             x(kk + 1 : kk + dof) = elem%zST(j,1:dof,k)

             kk = kk + dof
          enddo !j
        enddo !k

      enddo !i
      nullify( elem )

   end subroutine copyZST_toLongVector

   !> distributes  elem%zST to the global vector x to
   !> based on the global variables state%p_mod, state%q_mod
   !> puts zeros in the overfloating places in the vector
   subroutine copyZST_toLongVectorMod( grid, p_mod, q_mod, x )
      class( mesh ) , intent(in) :: grid
      integer, intent(in) :: p_mod, q_mod
      integer :: nsize
      real, dimension(:), intent(inout) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, kkUpdate

      nsize = state%bigNSize(p_mod, q_mod)
      if (nsize /= size(x) ) then
        print*, 'size of the vector = ', size(x) , 'but it should be ' , nsize

        stop 'wrong dimension of the output vector in copyZST_toLongVectorNew'
      endif

      if (state%getP_mod() /= p_mod) &
        print*, "copyZST_toLongVectorMod was not tested for state%getP_mod() /= p_mod"

      x(1:nsize) = 0.0

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (.not. allocated(elem%bigNcv) ) &
         stop 'not allocated bigNcv!'

        kk = elem%bigNcv(p_mod, q_mod)
        kkUpdate = DOFtriang( elem%deg + p_mod)

        do k = 1, elem%Tdof
          do j = 1, ndim
             x(kk : kk+dof-1) = elem%zST(j,1:dof,k)
             kk = kk + kkUpdate ! we have to get update according to p_mod ~ dof
          enddo !j
        enddo !k

      enddo !i
      nullify( elem )

   end subroutine copyZST_toLongVectorMod


   !> distributes  elem%zST to the global vector x to
   !> based on the global variables state%p_mod, state%q_mod
   !> puts zeros in the overfloating places in the vector
   subroutine copyZSTplus_toLongVectorMod( grid, p_mod, q_mod, x )
      class( mesh ) , intent(in) :: grid
      integer, intent(in) :: p_mod, q_mod
      integer :: nsize
      real, dimension(:), intent(inout) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, kkUpdate

      nsize = state%bigNSize(p_mod, q_mod)
      if (nsize /= size(x) ) then
        print*, 'size of the vector = ', size(x) , 'but it should be ' , nsize

        stop 'wrong dimension of the output vector in copyZST_toLongVectorNew'
      endif

      x(1:nsize) = 0.0

      do i = 1, grid%nelem
        elem => grid%elem(i)
        !dof = elem%dof

        if (.not. allocated(elem%bigNcv) ) &
         stop 'NOT allocated(elem%bigNcv)'

        kk = elem%bigNcv(p_mod, q_mod)
        kkUpdate = DOFtriang( elem%deg + p_mod)


        if (size(elem%zSTplus(1,:,1)) /= kkUpdate) &
         stop "wrong size of zSTplus in copyZSTplus_toLongVectorMod"

!        if (i==1) then
!            print*, "Control copyZSTplus_toLongVectorMod since only "
!            print*, "1:elem%dof from zSTplus is copied, its size", size(elem%zSTplus(1,:,1)), "dof = ", dof
!            print*, elem%zSTplus(1,:,1)
!            print*, "kkUpdate =" , kkUpdate
!        end if

        do k = 1, elem%Tdof
          do j = 1, ndim
             x(kk : kk+kkUpdate-1) = elem%zSTplus(j,1:kkUpdate,k)
             kk = kk + kkUpdate
          enddo !j
        enddo !k

     enddo !i
     nullify( elem )

   end subroutine copyZSTplus_toLongVectorMod

   !> distributes  elem%zST to the global vector x to
   !> based on the global variables state%p_mod, state%q_mod
   !> puts zeros in the overfloating places in the vector
   subroutine copyWST_toLongVectorMod( grid, p_mod, q_mod, x )
      class( mesh ) , intent(in) :: grid
      integer, intent(in) :: p_mod, q_mod
      integer :: nsize
      real, dimension(:), intent(inout) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof, kkUpdate

      nsize = state%bigNSize(p_mod, q_mod)
      if (nsize /= size(x) ) then
        print*, 'size of the vector = ', size(x) , 'but it should be ' , nsize

        stop 'wrong dimension of the output vector in copyZST_toLongVectorNew'
      endif

      x(1:nsize) = 0.0

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (.not. allocated(elem%bigNcv) ) &
         print*, 'elem%i', i

        kk = elem%bigNcv(p_mod, q_mod)
        kkUpdate = DOFtriang( elem%deg + p_mod)

        do k = 1, elem%Tdof
          do j = 1, ndim
             x(kk : kk+dof-1) = elem%wST(j,1:dof,k)
             kk = kk + kkUpdate
          enddo !j
        enddo !k

      enddo !i
      nullify( elem )

   end subroutine copyWST_toLongVectorMod

   !> distributes  elem%zST to the global vector x to
   !> based on the global variables state%p_mod, state%q_mod
   !> puts zeros in the overfloating places in the vector
   subroutine copyWSTplus_toLongVectorMod( grid, p_mod, q_mod, x )
      class( mesh ) , intent(in) :: grid
      integer, intent(in) :: p_mod, q_mod
      integer :: nsize
      real, dimension(:), intent(inout) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof, kkUpdate

      nsize = state%bigNSize(p_mod, q_mod)
      if (nsize /= size(x) ) then
        print*, 'size of the vector = ', size(x) , 'but it should be ' , nsize

        stop 'wrong dimension of the output vector in copyZST_toLongVectorNew'
      endif

      x(1:nsize) = 0.0

      print*, "Control copyWSTplus_toLongVectorMod since only "
      print*, "1:elem%dof from wSTplus is copied, its size", size(elem%wSTplus(1,:,1)), "dof = ", dof
      print*

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (.not. allocated(elem%bigNcv) ) &
         print*, 'elem%i', i

        kk = elem%bigNcv(p_mod, q_mod)
        kkUpdate = DOFtriang( elem%deg + p_mod)

        do k = 1, elem%Tdof
          do j = 1, ndim
             x(kk : kk+dof-1) = elem%wSTplus(j,1:dof,k)
             kk = kk + kkUpdate
          enddo !j
        enddo !k

      enddo !i
      nullify( elem )

   end subroutine copyWSTplus_toLongVectorMod

   !> distributes  elem%wST to the global vector x
   subroutine copyWST_toLongVector( grid,  nsize, x )
      class( mesh ) , intent(in) :: grid
      integer, intent(in) :: nsize
      real, dimension(1:nsize), intent(out) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof

      kk = 0

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (elem%ncv /= kk + 1) then
          print*, 'Problem while copying update from zST (solution.f90)'
          stop
        endif
        elemDof = elem%dof * ndim * elem%Tdof

        do k = 1, elem%Tdof
          do j = 1, ndim
             x(kk + 1 : kk + dof) = elem%wST(j,1:dof,k)
             kk = kk + dof
          enddo !j
        enddo !k

      enddo !i
      nullify( elem )

   end subroutine copyWST_toLongVector



   !> distributes the global solution vector x to elem%zST and elem%w(0,:)
   subroutine copyZST_fromLongVector( grid,  nsize, x, update )
      class( mesh ) , intent(inout) :: grid
      integer, intent(in) :: nsize
      real, dimension(1:nsize), intent(in) :: x
      logical, intent(in), optional :: update ! if .true. then zST = zST + x  else zST = x
      !      class( DWR_t ) , intent(in) :: DWR
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof

      kk = 0

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (elem%ncv /= kk + 1) then
          print*, 'Problem while copying update to wST (solution.f90)'
          stop
        endif
        elemDof = elem%dof * ndim * elem%Tdof

        if( present(update) .and. update ) then
           do k = 1, elem%Tdof
              do j = 1, ndim
                 !if(elem%i <= 2) &
                 !     write(*,'(a8, i5, 400es12.4)') 'ZST A:',elem%i, elem%zST(j,1:dof,k)
                 elem%zST(j,1:dof,k) = elem%zST(j,1:dof,k) + x(kk + 1 : kk + dof)
                 kk = kk + dof

                 !if(elem%i <= 2) then
                 !   write(*,'(a8, i5, 400es12.4)') 'ZST B:',elem%i, elem%zST(j,1:dof,k)
                 !   print*
                 !endif

              enddo !j
           enddo !k

        else
           do k = 1, elem%Tdof
              do j = 1, ndim
                 elem%zST(j,1:dof,k) = x(kk + 1 : kk + dof)
                 kk = kk + dof
              enddo !j
           enddo !k
        endif

        !  compute the solution in the endpoints

!        call Transfer_wST_to_w_Elem( elem , 0, elem%TQnum)
        elem%w(0, 1:dof*ndim) = Transfer_funST_to_fun( &
            elem%zST(1:ndim, 1:dof, 1:elem%Tdof), &
            dof, elem%Tdof, 0, elem%TQnum)
      enddo !i

      nullify( elem )

   end subroutine copyZST_fromLongVector

   !> distributes the global solution vector x to elem%zST and elem%w(0,:)
   subroutine copyZSTplus_fromLongVector( grid, nsize, p_mod, q_mod, x )
      class( mesh ) , intent(inout) :: grid
      integer, intent(in) :: nsize, p_mod, q_mod
      real, dimension(1:nsize), intent(in) :: x
!      class( DWR_t ) , intent(in) :: DWR
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof, Tdof

      kk = 0

      if (nsize /= state%bigNSize(p_mod,q_mod) ) &
         stop "wrong nsize in copyZSTplus_fromLongVectorMod"

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = DOFtriang( elem%deg + p_mod)
        Tdof = elem%Tdof + q_mod

        if (elem%bigNcv(p_mod,q_mod) /= kk + 1) then
          print*, 'Problem while copyZSTplus_fromLongVectorMod'
          stop
        endif
        elemDof = dof * ndim * Tdof

        do k = 1, Tdof
          do j = 1, ndim
             elem%zSTplus(j,1:dof,k) = x(kk + 1 : kk + dof)
             kk = kk + dof
          enddo !j
        enddo !k
        !  compute the solution in the endpoints
      enddo !i

      nullify( elem )

   end subroutine copyZSTplus_fromLongVector

   !> distributes the global solution vector x to elem%wST and elem%w(0,:)
   subroutine copyWST_fromLongVector( grid,  nsize, x )
      class( mesh ) , intent(inout) :: grid
      integer, intent(in) :: nsize
      real, dimension(1:nsize), intent(in) :: x
      class( element ), pointer :: elem
      integer :: i,j,k,kk, dof, elemDof

      kk = 0

      do i = 1, grid%nelem
        elem => grid%elem(i)
        dof = elem%dof

        if (elem%ncv /= kk + 1) then
          print*, 'Problem while copying update to wST (solution.f90)'
          stop
        endif
        elemDof = elem%dof * ndim * elem%Tdof

        do k = 1, elem%Tdof
          do j = 1, ndim
             elem%wST(j,1:dof,k) = x(kk + 1 : kk + dof)
             kk = kk + dof
          enddo !j
        enddo !k
        !  compute the solution in the endpoints

!        call Transfer_wST_to_w_Elem( elem , 0, elem%TQnum)
        elem%w(0, 1:dof*ndim) = Transfer_funST_to_fun( &
            elem%wST(1:ndim, 1:dof, 1:elem%Tdof), &
            dof, elem%Tdof, 0, elem%TQnum)

      enddo !i

      nullify( elem )

   end subroutine copyWST_fromLongVector

   !copy elem%w(0,:) to wS
   subroutine transfer_w_to_WS_elem( elem )
     class( element ), intent(inout) :: elem
     integer :: j, dof

     dof = elem%dof

     if ( allocated( elem%wS ) ) then
         deallocate( elem%wS )
     endif

     allocate( elem%wS(1:ndim, 1:elem%dof) )

     do j = 1, ndim
         elem%wS(j, 1:dof) = elem%w( 0, (j-1)*dof + 1 : j*dof )
     enddo !j

   end subroutine transfer_w_to_WS_elem


   !> compute (pseudo maximal norm) - in space integration nodes
   !> in the time interval endpoint
  function solution_L8norm(grid) result(norm)
   class( mesh ), intent(in) :: grid
   real :: norm

   class( element ), pointer :: elem
   real, allocatable, dimension(:,:)  :: wi
   real :: temp
   integer :: i

      norm = 0.0

      do i = 1,grid%nelem
         elem => grid%elem(i)
         allocate( wi(1:elem%Qdof, 1:ndim), source = 0.0 )

         call Eval_whST_Elem( elem, 1, 0, wi)
         temp = maxval( abs( wi(:,:) ) )
         norm = max( norm, temp)
         deallocate( wi )
      end do
  end function solution_L8norm

!  !> copy array elem(:)%wST to zST and
!  !> put ZEROS into wST
!  subroutine copyWSTtoZST( grid )
!    class( mesh ), intent(inout) :: grid
!    class( element ), pointer :: elem
!    integer :: i
!
!    ! false - not deg_plus
!    call allocateZST( grid, .false. )
!
!    do i = 1, grid%nelem
!      elem => grid%elem(i)
!
!      elem%zST(1:ndim, 1:elem%dof, 1:elem%Tdof) = elem%wST( 1:ndim, 1:elem%dof, 1:elem%Tdof)
!      elem%wST(1:ndim, 1:elem%dof, 1:elem%Tdof) = 0.0
!
!    end do !
!
!  end subroutine copyWSTtoZST

!  subroutine moveZSTtoZSTplus( grid )
!    class( mesh ), intent(inout) :: grid
!    class( element ), pointer :: elem
!    integer :: i
!
!    print*, 'move ZST to ZSTplus called'
!
!    do i = 1,grid%nelem
!
!      elem => grid%elem(i)
!      if ( associated(elem%zST) ) then
!         if ( associated(elem%zSTplus) ) then
!          if ( associated(elem%zSTplus, target= elem%zST) ) then
!            print*, 'second time called '
!          else
!            deallocate( elem%zSTplus)
!          endif
!         endif
!
!         elem%zSTplus => elem%zST
!         nullify( elem%zST )
!         ! projection of z_{p+1} to z_p
!         allocate( elem%zST(1:ndim, 1:elem%dof, 1:elem%Tdof) , source = 0.0 )
!
!         ! due to the OG of basis function this is the projection prom P^{p+1} to P^p
!         !FERROR : WOULD NOT WORK FOR CURVED ELEMENTS
!         elem%zST(1:ndim, 1:elem%dof, 1:elem%Tdof) = &
!               elem%zSTplus(1:ndim, 1:elem%dof, 1:elem%Tdof)
!      else
!        stop 'zST is not associated in moveZSTtoZSTplus'
!      end if
!
!    end do !i
!
!  end subroutine moveZSTtoZSTplus

!  !> variant of subroutine, which does not deallocate elem%zSTplus arrays
!  subroutine moveZSTtoZSTplus_simple( grid )
!    class( mesh ), intent(inout) :: grid
!    class( element ), pointer :: elem
!    integer :: i, dofP
!
!    print*, 'move ZST to ZSTplus SIMPLE called'
!
!    do i = 1,grid%nelem
!
!      elem => grid%elem(i)
!      if ( associated(elem%zST) ) then
!         !if ( associated(elem%zSTplus) ) then
!         ! if ( associated(elem%zSTplus, target= elem%zST) ) then
!         !   print*, 'second time called '
!         ! else
!         !   deallocate( elem%zSTplus)
!         ! endif
!         !endif
!
!         dofP = size (elem%zST, 2)
!
!         elem%zSTplus(1:ndim, :, -1) = 0.
!         elem%zSTplus(1:ndim, 1:dofP, -1) =  elem%zST(1:ndim, 1:dofP, 1)
!
!
!
!         nullify( elem%zST )
!         ! projection of z_{p+1} to z_p
!         allocate( elem%zST(1:ndim, 1:elem%dof, 1:elem%Tdof) , source = 0.0 )
!
!         ! due to the OG of basis function this is the projection prom P^{p+1} to P^p
!         !FERROR : WOULD NOT WORK FOR CURVED ELEMENTS
!         elem%zST(1:ndim, 1:elem%dof, 1) = elem%zSTplus(1:ndim, 1:elem%dof, -1)
!      else
!        stop 'zST is not associated in moveZSTtoZSTplus_simple'
!      end if
!
!    end do !i
!
!  end subroutine moveZSTtoZSTplus_simple

end module solution_mod
