module bmproblem

  use iso_c_binding, only: c_ptr,c_f_pointer
  
  implicit none
  
  type :: pdata_type
     integer :: np,nprob
     real(kind=8) :: h,theta
     integer, allocatable :: ipar(:)
     real(kind=8), allocatable :: phi(:,:,:),uopt(:,:,:),phi2D(:,:),uopt2D(:,:),rpar(:),x0(:),A(:,:),b(:)
  end type pdata_type

  public :: iniprob,endprob,evalr

contains

  ! ******************************************************************
  ! ******************************************************************

  subroutine iniprob(n,pdataptr)

    implicit none

    ! SCALAR ARGUMENTS
    integer, intent(out) :: n
    type(c_ptr), optional, intent(in) :: pdataptr
  
    ! LOCAL SCALARS
    integer :: allocstat,i,j,k,lrpar
    real(kind=8) :: a,b,bll,bur,c,d,h,lambda,seed,src,tmp,Re
    type(pdata_type), pointer :: pdata

    ! EXTERNAL FUNCTIONS
    real(kind=8), external :: drand
    
    call c_f_pointer(pdataptr,pdata)

    write(*,*) 'Enter problem number: '
    read(*,*) pdata%nprob
    
    write(*,*) 'Number of points (np): '
    read(*,*) pdata%np

    if ( pdata%nprob .eq. 21 ) then
       ! THIS BRATU 2D
       
       n = ( pdata%np - 2 ) ** 2

       pdata%h = 1.0d0 / ( pdata%np - 1 )

       write(*,*) 'Enter theta: '
       read(*,*) pdata%theta
       
       allocate(pdata%uopt2D(pdata%np,pdata%np),pdata%phi2D(2:pdata%np-1,2:pdata%np-1),pdata%x0(n),stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if
    
       ! Initial guess

       pdata%x0(1:n) = 0.0d0
       
       ! "Known solution" used to generate the problem data

       do k = 1,pdata%np
          do j = 1,pdata%np
             a = real( pdata%np - k ) / ( pdata%np - 1 )
             b = real( pdata%np - j ) / ( pdata%np - 1 )
             pdata%uopt2D(j,k) = 10.0d0 * a * b * (1.0d0 - a) * (1.0d0 - b) * exp( a ** 4.5d0 )
          end do
       end do
       
       ! Right hand side of the system

       tmp = 1.0d0 / pdata%h ** 2
    
       do k = 2,pdata%np - 1
          do j = 2,pdata%np - 1
             pdata%phi2D(j,k) = - tmp * ( pdata%uopt2D(j,k+1) + pdata%uopt2D(j,k-1) + pdata%uopt2D(j+1,k) + &
                  pdata%uopt2D(j-1,k) - 4.0d0 * pdata%uopt2D(j,k) ) + pdata%theta * exp( pdata%uopt2D(j,k) )
          end do
       end do
       
    else if ( pdata%nprob .eq. 22 ) then
       ! THIS IS BRATU 3D

       n = ( pdata%np - 2 ) ** 3

       pdata%h = 1.0d0 / ( pdata%np - 1 )

       write(*,*) 'Enter theta: '
       read(*,*) pdata%theta
       
       allocate(pdata%uopt(pdata%np,pdata%np,pdata%np), &
            pdata%phi(2:pdata%np-1,2:pdata%np-1,2:pdata%np-1),pdata%x0(n),stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if
    
       ! Initial guess

       pdata%x0(1:n) = 0.0d0
       
       ! "Known solution" used to generate the problem data

       do k = 1,pdata%np
          do j = 1,pdata%np
             do i = 1,pdata%np
                a = real( pdata%np - k ) / ( pdata%np - 1 )
                b = real( pdata%np - j ) / ( pdata%np - 1 )
                c = real( pdata%np - i ) / ( pdata%np - 1 )
                pdata%uopt(i,j,k) = 10.0d0 * a * b * c &
                     * (1.0d0 - a) * (1.0d0 - b) * (1.0d0 - c) * exp( a ** 4.5d0 )
             end do
          end do
       end do
       
       ! Right hand side of the system

       tmp = 1.0d0 / pdata%h ** 2
    
       do k = 2,pdata%np - 1
          do j = 2,pdata%np - 1
             do i = 2,pdata%np - 1
                pdata%phi(i,j,k) = - tmp * ( pdata%uopt(i,j,k+1) + pdata%uopt(i,j,k-1) + pdata%uopt(i,j+1,k) + &
                     pdata%uopt(i,j-1,k) + pdata%uopt(i+1,j,k) + pdata%uopt(i-1,j,k) - 6.0d0 * pdata%uopt(i,j,k) ) + &
                     pdata%theta * exp( pdata%uopt(i,j,k) )
             end do
          end do
       end do
       
    else if ( pdata%nprob .eq. 23 ) then
       ! This is Cavity 2D from NITSOL distribution
       
       n = pdata%np ** 2
       
       lrpar = 100 * ( pdata%np ** 2 + pdata%np + 1 )
      
       allocate(pdata%ipar(2),pdata%rpar(lrpar),pdata%x0(n),stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if

       pdata%ipar(1) = pdata%np
       pdata%ipar(2) = pdata%np + 2

       Re = 1.0d0 ! Suggested in NITSOL paper (section 4.3) 
       
       call init(n,pdata%x0,Re,pdata%rpar,pdata%ipar)

    else if ( pdata%nprob .eq. 24 ) then
       ! This is Generalized Bratu 2D from NITSOL distribution
       
       n = pdata%np ** 2
       
       d = 32.0d0
       lambda = 16.0d0

       allocate(pdata%ipar(2),pdata%rpar(3),pdata%x0(n),stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if

       pdata%ipar(1) = pdata%np
       pdata%ipar(2) = pdata%np

       h = 1.0d0 / real( pdata%np + 1 )
       
       pdata%rpar(1) = 1.0d0 - h * d / 2.0d0
       pdata%rpar(2) = 1.0d0 + h * d / 2.0d0
       pdata%rpar(3) = h * h * lambda

       pdata%x0(1:n) = 0.0d0

    else if ( pdata%nprob .eq. 25 ) then
       ! This is Pormed 2D from NITSOL distribution
       
       n = pdata%np ** 2

       d = 50.0d0 
       bll = 1.0d0
       bur = 0.0d0
       src = 10.0d0
       
       allocate(pdata%ipar(2),pdata%rpar(4),pdata%x0(n),stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if

       pdata%ipar(1) = pdata%np
       pdata%ipar(2) = pdata%np

       h = 1.0d0 / real( pdata%np + 1 )
       
       pdata%rpar(1) = bll
       pdata%rpar(2) = bur
       pdata%rpar(3) = h * h * src
       pdata%rpar(4) = h * d / 2.0d0 

       do j = 1,pdata%np
          do i = 1,pdata%np
             k = ( j - 1 ) * pdata%np + i
             pdata%x0(k) = 1.0d0 - real( i ) * real( j ) * h * h
          end do
       end do

    else if ( pdata%nprob .eq. 26 ) then
       ! This is a linear system for testing
       n = pdata%np

       allocate(pdata%A(n,n),pdata%b(n),pdata%x0(n),stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if
       
       pdata%x0(1:n) = 0.0d0

       seed = 123456.0d0
       do j = 1,n
          do i = 1,n
             pdata%A(i,j) = 2.0d0 * drand(seed) - 1.0d0
          end do
       end do

       pdata%b(1:n) = 0.0d0
       do j = 1,n
          pdata%b(1:n) = pdata%b(1:n) + pdata%A(1:n,j)
       end do
    end if
  
  end subroutine iniprob

  ! ******************************************************************
  ! ******************************************************************

  subroutine endprob(pdataptr)

    implicit none

    ! SCALAR ARGUMENTS
    type(c_ptr), optional, intent(in) :: pdataptr
  
    ! LOCAL SCALARS
    integer :: allocstat
    type(pdata_type), pointer :: pdata
    
    call c_f_pointer(pdataptr,pdata)
  
    if ( pdata%nprob .eq. 21 ) then
       deallocate(pdata%uopt2D,pdata%phi2D,pdata%x0,stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if

    else if ( pdata%nprob .eq. 22 ) then
       deallocate(pdata%uopt,pdata%phi,pdata%x0,stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if

    else if ( pdata%nprob .eq. 23 .or. pdata%nprob .eq. 24 .or. pdata%nprob .eq. 25 ) then
       deallocate(pdata%ipar,pdata%rpar,pdata%x0,stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if

    else if ( pdata%nprob .eq. 26 ) then
       deallocate(pdata%A,pdata%b,pdata%x0,stat=allocstat)
       if ( allocstat .ne. 0 ) then
          write(*,*) 'Error in allocate statement!'
          stop
       end if
    end if
    
  end subroutine endprob
  
  ! ******************************************************************
  ! ******************************************************************

  subroutine evalr(n,x,r,pdataptr)

    implicit none

    ! SCALAR ARGUMENTS
    integer, intent(in) :: n
    type(c_ptr), optional, intent(in) :: pdataptr

    ! ARRAY ARGUMENTS
    real(kind=8), intent(in) :: x(n)
    real(kind=8), intent(out) :: r(n)
    
    ! LOCAL SCALARS
    integer :: flag,i,j,k,ell
    real(kind=8) :: tmp,xi0,xi1,xi2,xi3,xi4,xi5,xi6
    type(pdata_type), pointer :: pdata
    
    call c_f_pointer(pdataptr,pdata)
    
    if ( pdata%nprob .eq. 21 ) then
       tmp = 1.0d0 / pdata%h ** 2
       
       ell = 0
       do j = 2,pdata%np - 1
          do k = 2,pdata%np - 1
             ell = ell + 1
                
             xi0 = x(xind2D(j,k,pdata%np))
              
             if ( k .eq. pdata%np - 1 ) then
                xi1 = pdata%uopt2D(j,k+1)
             else
                xi1 = x(xind2D(j,k+1,pdata%np))
             end if
              
             if ( k .eq. 2 ) then
                xi2 = pdata%uopt2D(j,k-1)
             else
                xi2 = x(xind2D(j,k-1,pdata%np))
             end if
              
             if ( j .eq. pdata%np - 1 ) then
                xi3 = pdata%uopt2D(j+1,k)
             else
                xi3 = x(xind2D(j+1,k,pdata%np))
             end if
              
             if ( j .eq. 2 ) then
                xi4 = pdata%uopt2D(j-1,k)
             else
                xi4 = x(xind2D(j-1,k,pdata%np))
             end if
              
             r(ell) = - tmp * ( xi1 + xi2 + xi3 + xi4 - 4.0d0 * xi0 ) + pdata%theta * exp( xi0 ) - pdata%phi2D(j,k)
          end do
       end do

    else if ( pdata%nprob .eq. 22 ) then
       tmp = 1.0d0 / pdata%h ** 2
       
       ell = 0
       do i = 2,pdata%np - 1
          do j = 2,pdata%np - 1
             do k = 2,pdata%np - 1
                ell = ell + 1
                
                xi0 = x(xind(i,j,k,pdata%np))
              
                if ( k .eq. pdata%np - 1 ) then
                   xi1 = pdata%uopt(i,j,k+1)
                else
                   xi1 = x(xind(i,j,k+1,pdata%np))
                end if
              
                if ( k .eq. 2 ) then
                   xi2 = pdata%uopt(i,j,k-1)
                else
                   xi2 = x(xind(i,j,k-1,pdata%np))
                end if
              
                if ( j .eq. pdata%np - 1 ) then
                   xi3 = pdata%uopt(i,j+1,k)
                else
                   xi3 = x(xind(i,j+1,k,pdata%np))
                end if
              
                if ( j .eq. 2 ) then
                   xi4 = pdata%uopt(i,j-1,k)
                else
                   xi4 = x(xind(i,j-1,k,pdata%np))
                end if
              
                if ( i .eq. pdata%np - 1 ) then
                   xi5 = pdata%uopt(i+1,j,k)
                else
                   xi5 = x(xind(i+1,j,k,pdata%np))
                end if
              
                if ( i .eq. 2 ) then
                   xi6 = pdata%uopt(i-1,j,k)
                else
                   xi6 = x(xind(i-1,j,k,pdata%np))
                end if
              
                r(ell) = - tmp * ( xi1 + xi2 + xi3 + xi4 + xi5 + xi6 - 6.0d0 * xi0 ) + pdata%theta * exp( xi0 ) - pdata%phi(i,j,k)
             end do
          end do
       end do

    else if ( pdata%nprob .eq. 23 ) then
       call fcav(n,x,r,pdata%rpar,pdata%ipar,flag)

    else if ( pdata%nprob .eq. 24 ) then
       call fbratu(n,x,r,pdata%rpar,pdata%ipar,flag)

    else if ( pdata%nprob .eq. 25 ) then
       call fpm(n,x,r,pdata%rpar,pdata%ipar,flag)

    else if ( pdata%nprob .eq. 26 ) then
       r(1:n) = matmul( pdata%A(1:n,1:n), x(1:n) ) - pdata%b(1:n)
    end if
    
  end subroutine evalr

  ! ******************************************************************
  ! ******************************************************************
  
  integer function xind2D(j,k,np)
    
    implicit none
    
    ! SCALAR ARGUMENTS
    integer, intent(in) :: j,k,np
    
    xind2D = ( np - 2 ) * ( j - 2 ) + k - 1
    
  end function xind2D

  ! ******************************************************************
  ! ******************************************************************
  
  integer function xind(i,j,k,np)
    
    implicit none
    
    ! SCALAR ARGUMENTS
    integer, intent(in) :: i,j,k,np
    
    xind = ( np - 2 ) ** 2 * ( i - 2 ) + ( np - 2 ) * ( j - 2 ) + k - 1
    
  end function xind

end module bmproblem
