! pde_abc_eq.f90  discretization using module abc
! solve a1(x,y) uxx(x,y) + b1(x,y) uxy(x,y) + c1(x,y) uyy(x,y) +
!       d1(x,y) ux(x,y)  + e1(x,y) uy(x,y)  + f1(x,y) u(x,y) = c(x,y)
! functions and other problem information are in the file  abc.f90
! method is high order discretization
! boundary conditions computed using u(x)
! analytic solution may be given by u(x)
! solve for Uij numerical approximation U(x_i,y_j) of u(x_i,y_j)
! ut * us = uy, solve simultaneous equations for U
!
! precompile module abc.f90, compile with deriv.f90 and simeq.f90

module global
  use abc  ! defines nx, ny, np, ifcheck, xmin..., a1, b1, ...
  implicit none
  integer, parameter :: nxy = (nx-2)*(ny-2)
  double precision, dimension(nxy,nxy) :: ut
                      ! ((i-1)*(ny-1)+(ii-1),(j-1)*(ny-1)+(jj-1))
  double precision, dimension(nxy) :: uy ! RHS, some times in ut
  double precision, dimension(nxy) :: us ! solution being computed        
  double precision, dimension(nxy) :: uu ! solution if checking
  double precision, dimension(nx) :: xg    ! X coordinates
  double precision, dimension(ny) :: yg    ! Y coordinates
  double precision :: val, err, error, avgerr, maxerr
  double precision :: hx, hy, hx2, hy2, hxy
  double precision, dimension(nd,nd) :: pdx
  double precision, dimension(nd,nd) :: pdy
  double precision, dimension(nd,nd) :: pd
  double precision, dimension(nd,nd) :: pdxx
  double precision, dimension(nd,nd) :: pdyy
  double precision, dimension(nd,nd) :: pdxy
  integer cs, md

  interface
    subroutine simeq(n, sz, A, Y, X)
      integer, intent(in) :: n  ! number of equations
      integer, intent(in) :: sz ! dimension of arrays
      double precision, dimension(sz,sz), intent(in) :: A
      double precision, dimension(sz), intent(in) :: Y
      double precision, dimension(sz), intent(out) :: X
    end subroutine simeq
    subroutine deriv(order, npoints, point, a, bb)
      integer, intent(in) :: order    ! derivative order
      integer, intent(in) :: npoints  ! number of points to be used
      integer, intent(in) :: point    ! compute derivative at this point
      integer, intent(out), dimension(*) :: a ! returned coefficients
      integer, intent(out) :: bb      ! divisor, also use hx^order
    end subroutine deriv 
    subroutine rderiv(order, npoints, point, h, aa)
      integer, intent(in) :: order    ! derivative order
      integer, intent(in) :: npoints  ! number of points to be used
      integer, intent(in) :: point    ! compute derivative at this point
      double precision, intent(in) :: h  ! divisor for use h^order
      double precision, intent(out), dimension(*)::aa ! returned coefficients
    end subroutine rderiv 
    function S(i,j) result(k)
      integer, intent(in) :: i, j
      integer :: k
    end function S
  end interface
end module global

function S(i,j) result(k) ! subscripting function, generalize to higher dim
  use abc
  implicit none
  integer, intent(in) :: i, j
  integer :: k
  k = (i-2)*(ny-2)+(j-1)
end function S

subroutine clear()
  use global
  implicit none
  integer :: i, j, ii

  do i=2,nx-1
    do j = 2,ny-1
      ii = S(i,j)
       us(ii) = 0.0   ! inside boundary
       if(ifcheck>0) uu(ii) = u(xg(i),yg(j))
    end do
  end do
end subroutine clear

subroutine printus()
  use global
  implicit none
  integer :: i, j
  
  print "(/a)", "computed solution"
  print *, "us(x,y)"
  do i=2,nx-1
    print "(11f7.3)", (us(S(i,j)), j=2,min(ny-1,11))
  end do
  print *, " "
end subroutine printus 

subroutine check() ! typically not known, instrumentation 
  use global
  implicit none
  integer :: i, j
  double precision :: uexact
  
  error = 0.0
  maxerr = 0.0
  do i=2,nx-1
    do j=2,ny-1
        uexact = u(xg(i),yg(j))
        err = abs(us(S(i,j))-uexact)
        error = error + err
        if(err>maxerr) maxerr=err
    end do
  end do
end subroutine check 

subroutine printexact()
  use global
  implicit none
  integer :: i, j
  
  print "(/a)", "exact solution"
  print *, "u(x,y)"
  do i=2,nx-1
    print "(11f7.3)", (u(xg(i),yg(j)),j=2,min(ny-1,11))
  end do
  print *, " "
end subroutine printexact 

subroutine printdiff()
  use global
  implicit none
  integer :: i, j
  
  print "(/a)", "difference exact-computed solution"
  print *, "u(x,y)-us(x,y)" 
  do i=2,nx-1
    print "(10f8.4)", (u(xg(i),yg(j))-us(S(i,j)), j=2,min(ny-1,10))
  end do
  print *, " "
end subroutine printdiff 

subroutine print_coef()
  use global
  implicit none
  integer :: i, j

  print *, "pdxx    j=1     j=2     j=3     j=4     j=5  y"
  do i=1,nd
    print "('i=',I1,',',9f8.2,' ')", i, (pdxx(i,j), j=1,min(nd,9))
  end do
  print *, "x"
  print *, " "
  print *, "pdxy    j=1     j=2     j=3     j=4     j=5  y"
  do i=1,nd
    print "('i=',I1,',',9f8.2,' ')", i, (pdxy(i,j), j=1,min(nd,9))
  end do
  print *, "x"
  print *, " "
  print *, "pdyy    j=1     j=2     j=3     j=4     j=5  y"
  do i=1,nd
    print "('i=',I1,',',9f8.2,' ')", i, (pdyy(i,j), j=1,min(nd,9))
  end do
  print *, "x"
  print *, " "
  print *, "pdx     j=1     j=2     j=3     j=4     j=5  y"
  do i=1,nd
    print "('i=',I1,',',9f8.2,' ')", i, (pdx(i,j), j=1,min(nd,9))
  end do
  print *, "x"
  print *, " "
  print *, "pdy     j=1     j=2     j=3     j=4     j=5  y"
  do i=1,nd
    print "('i=',I1,',',9f8.2,' ')", i, (pdy(i,j), j=1,min(nd,9))
  end do
  print *, "x"
  print *, " "
end subroutine print_coef

subroutine build_coef(ix, iy, ic)
  use global
  implicit none
  integer, intent(in) :: ix, iy, ic
  integer :: i, j, k, bb
  integer, dimension(25) :: a

  if(((nd/2)*2.eq.nd) .or. nd.lt.5) then
     print *, "nd must be odd and >1 , bye."
     stop
  end if

  ! build discretization matrices
  if(ifcheck>0) print "('build_coef for point ix=',I2,', iy=',I2)", ix, iy
  do i=1,nd
    do j=1,nd
      pdxx(i,j)=0.0 ! coefficient of a1(x,y)
      pdxy(i,j)=0.0 ! coefficient of b1(x,y)
      pdyy(i,j)=0.0 ! coefficient of c1(x,y)
      pdx (i,j)=0.0 ! coefficient of d1(x,y)
      pdy (i,j)=0.0 ! coefficient of e1(x,y)
      pd  (i,j)=0.0 ! coefficient of f1(x,y)
    end do
  end do
  pd(ix,iy) = 1.0
  call deriv(2, nd, ix-1, a, bb)  ! for pdxx, zero based point
  if(ifcheck>3) print "('deriv pdxx returns bb=',I2,', aa=',10I5)", &
                     bb, (a(k),k=1,min(nd,10))
  do i=1,nd
    pdxx(i,iy) = a(i)/(bb*hx2)
  end do

  call deriv(2, nd, iy-1, a, bb)  ! for pdyy 
  if(ifcheck>3) print "('deriv pdyy returns bb=',I2,', aa=',10I5)", &
                     bb, (a(k),k=1,min(nd,10))
  do j=1,nd
    pdyy(ix,j) = a(j)/(bb*hy2)
  end do

  call deriv(1, nd, ix-1, a, bb) ! for pdx 
  if(ifcheck>3) print "('deriv pdx returns bb=',I2,', aa=',10I5)", &
                     bb, (a(k),k=1,min(nd,10))
  do i=1,nd
    pdx(i,iy) = a(i)/(bb*hx)
  end do

  call deriv(1, nd, iy-1, a, bb) ! for pdy 
  if(ifcheck>3) print "('deriv pdy returns bb=',I2,', aa=',10I5)", &
                     bb, (a(k),k=1,min(nd,10))
  do j=1,nd
    pdy(ix,j) = a(j)/(bb*hy)
  end do

  do i=1,nd
    do j=1,nd
      pdxy(i,j) = pdx(i,iy)*pdy(ix,j)
    end do
  end do
  if(ifcheck>2) call print_coef()

end subroutine build_coef

subroutine check_row(row)
  use global
  implicit none
  integer, intent(in) :: row
  integer :: i, j, ii
  double precision :: sum

  
  sum = 0.0;
  do i=2,nx-1
    do j=2,ny-1
      ii = S(i,j)
      sum = sum + ut(row,ii)*uu(ii);
    end do
  end do
  sum = sum - uy(row)
  if(abs(sum)>0.001) print "('row=',I2,' is bad err=',E15.7)", row, sum
end subroutine check_row

subroutine build_row(i, j, bi, bj)
  use global
  implicit none
  integer, intent(in) :: i, j, bi, bj
  integer :: ki, kj, ii, jj, ij, ir, jr, ia, ja
  double precision :: p, a1c, b1c, c1c, d1c, e1c, f1c, cc


  ir = i
  jr = j
  ii = S(i,j) ! row of matrix 
  if(ifcheck>2) then
    print "('working row ',I2,', i=',I2,', j=',I2,', bi=',I2,', bj=',I2)", &
           ii, i, j, bi, bj
  end if
  a1c = a1(xg(ir),yg(jr))
  b1c = b1(xg(ir),yg(jr))
  c1c = c1(xg(ir),yg(jr))
  d1c = d1(xg(ir),yg(jr))
  e1c = e1(xg(ir),yg(jr))
  f1c = f1(xg(ir),yg(jr))
  cc  =  c(xg(ir),yg(jr))
  do ki=1,nd
    do kj=1,nd
      ia = ir-bi+ki-md ! real i,j in all element coordinates 
      ja = jr-bj+kj-md ! may be boundary 
      ij = S(ia,ja)    ! in ut coordinates 
      p = a1c*pdxx(ki,kj) &
        + b1c*pdxy(ki,kj) &
        + c1c*pdyy(ki,kj) &
        + d1c* pdx(ki,kj) &
        + e1c* pdy(ki,kj) &
        + f1c*  pd(ki,kj)
      if(ia<1)  print "('***-error in ia=',I2)", ia
      if(ia>nx) print "('***+error in ia=',I2)", ia
      if(ja<1)  print "('***-error in ja=',I2)", ja
      if(ja>ny) print "('***+error in ja=',I2)", ja
      if(ia.eq.1) then ! boundary value tests 
        uy(ii) = uy(ii) - p * u(xg(ia), yg(ja))
      elseif(ia.eq.nx) then
        uy(ii) = uy(ii) - p * u(xg(ia), yg(ja))
      elseif(ja.eq.1) then
        uy(ii) = uy(ii) - p * u(xg(ia), yg(ja))
      elseif(ja.eq.ny) then
        uy(ii) = uy(ii) - p * u(xg(ia), yg(ja))
      else ! non boundary case 
        ut(ii,ij) = p;
        ! if(ifcheck>0.and.((i<=3 .and. j<=3).or.(i>=nx-3 .and. j>=ny-3))) &
        !     print "('ij=',I2,', u(ii,ij)=',f9.3)", ij, ut(ii,ij)
      end if
    end do
  end do
  uy(ii) = uy(ii) + cc
  print "('ii=',I2,', ut(ii,ii)=',f8.3,', uy(ii)=',f8.3)", &
         ii, ut(ii,ii), uy(ii)
  if(ifcheck>0) call check_row(ii)
end subroutine build_row


subroutine init_matrix()
  use global
  implicit none
  integer :: i, j, ii, jj
  
  ! zero internal cells 
  do i=2,nx-1
    do j=2,ny-1
      do ii=2,nx-1
        do jj=2,ny-1 
          ut(S(i,j),S(ii,jj)) = 0.0
        end do
      end do
    end do
  end do

  print *, "internal cells zeroed "          
  call build_coef(md, md, md)
  do i=3,nx-2 ! compute standard central cells, real coordinates
    do j=3,ny-2
      call build_row(i, j, 0, 0)
    end do
  end do
  print *, 'standard central cells generated'

  ! do four sides, one element in. special cases
  call build_coef(md-1, md, md)
  do j=3,ny-2
    i = 2
    if(ifcheck>2) print "('working row edge i=',I2,', j=',I2)", i, j
    call build_row(i, j, -1, 0)
  end do ! j loop on edge

  call build_coef(md+1, md, md)
  do j=3,ny-2
    i = nx-1
    if(ifcheck>2) print "('working row edge i=',I2,', j=',I2)", i, j
    call build_row(i, j, 1, 0)
  end do ! j loop on edge

  call build_coef(md, md-1, md)
  do i=3,nx-2
    j = 2
    if(ifcheck>2) print "('working row edge i=',I2,', j=',I2)", i, j
    call build_row(i, j, 0 , -1)
  end do ! i loop on edge

  call build_coef(md, md+1, md)
  do i=3,nx-2
    j = ny-1
    if(ifcheck>2) print "('working row edge i=',I2,', j=',I2)", i, j
    call build_row(i, j, 0 , 1)
  end do ! i loop on edge

  print *,"four sides, one element in, special case, generated "
  ! now do four corners
  i=2
  j=2
  call build_coef(md-1, md-1, md)
  call build_row(i, j, -1, -1)
  i=2
  j=ny-1
  call build_coef(md-1, md+1, md)
  call build_row(i, j, -1, 1)
  i=nx-1
  j=2
  call build_coef(md+1, md-1, md)
  call build_row(i, j, 1, -1)
  i=nx-1
  j=ny-1
  call build_coef(md+1, md+1, md);
  call build_row(i, j, 1, 1);
end subroutine init_matrix
  
subroutine print_mat()
  use global
  implicit none
  integer :: i, ii, j, jj 
  print "(/a)", "initial matrix"
  do i=2,nx-1
    do ii=2,ny-1
      print "('i=',I3,', j=',I3,' .. ',I3)", S(i,ii),S(2,2),S(nx-1,ny-1)
      do j=2,nx-1
        print "(8f9.3)", (ut(S(i,ii),S(j,jj)),jj=2,min(ny-1,8))
      end do
      print "(f9.3)", uy(S(i,ii))
      print *, " "
    end do
  end do
  print *, " "
end subroutine print_mat

subroutine check_soln() ! against PDE, added later, used rderiv and xg,yg 
               ! for all interior points 
  use global
  implicit none
  double precision :: x, y, rmserr
  double precision :: sumerr, sumsqerr
  double precision, dimension(nx,nx) :: cx
  double precision, dimension(nx,nx) :: cxx
  double precision, dimension(ny,ny) :: cy
  double precision, dimension(ny,ny) :: cyy
  double precision, dimension(max(nx,ny)) :: tcxy
  double precision :: dxx, dxy, dyy, dx, dy ! a1, b1, c1, d1, e1 coefficients 
  double precision, dimension(nx,ny) :: ug ! full, including boundaries,
                                           ! copied from us() 
  integer :: cnt
  integer i, ii, j, jj, k

  cnt = 0
  if(ifcheck>5) print *, "check_soln setup"
  do i=1,nx
    xg(i) = xmin + (i-1)*hx
  end do
  do j=1,ny
    yg(j) = ymin + (j-1)*hy
  end do
  ! transform us() into ug, add boundaries, make code simpler 
  do i=2,nx-1 ! get from us() 
    do j=2,ny-1
      ii = S(i,j) ! ii for  us(ii) computed solution at i,j 
      ug(i,j) = us(ii)
      if(ifcheck>5) then
        print "('ug(',I2,',',I2,')=',e15.6)", i, j, ug(i,j)
      end if
    end do
  end do
  do i=1,nx ! fill boundaries, loops below easier 
    ug(i,1)  = u(xg(i),yg(1))
    ug(i,ny) = u(xg(i),yg(ny))
  end do
  do j=1,ny ! fill boundaries 
    ug(1,j)  = u(xg(1), yg(j))
    ug(nx,j) = u(xg(nx),yg(j))
  end do

  do i=1,nx
    call rderiv(1,nx,i-1,hx,tcxy) ! fills row of derivative matrix
    do k=1,nx
      cx(i,k) = tcxy(k)
    end do 
    call rderiv(2,nx,i-1,hx,tcxy)
    do k=1,nx
      cxx(i,k) = tcxy(k)
    end do 
  end do
  do j=1,ny
    call rderiv(1,ny,j-1,hy,tcxy)
    do k=1,ny
      cy(j,k) = tcxy(k)
    end do 
    call rderiv(2,ny,j-1,hy,tcxy)
    do k=1,ny
      cyy(j,k) = tcxy(k)
    end do 
  end do
  if(ifcheck>5) then
    print *, "cx="
    do i=1,nx
      do k=1,nx
        print "(f8.3)",cx(i,k)
      end do
      print *, " "
    end do
    print *, " "
    print *, "cxx="
    do i=1,nx
      do k=1,nx
        print "(f8.3)",cxx(i,k)
      end do
      print *, " "
    end do
    print *, " "
    print *, "cy="
    do j=1,ny
      do k=1,ny
        print "(f8.3)",cy(j,k)
      end do
      print *, " "
    end do
    print *, " "
    print *, "cyy="
    do j=1,ny
      do k=1,ny
        print "(f8.3)",cyy(j,k)
      end do
      print *, " "
    end do
    print *, " "
  end if ! ifcheck 

  sumerr = 0.0
  sumsqerr = 0.0
  maxerr = 0.0
  do i=2,nx-1 ! check full equation, at every point 
    x = xg(i)
    do j=2,ny-1
      y = yg(j)
      ! compute derivatives, plug into PDE, subtract RHS, should be zero 
      dxx = 0.0
      dx  = 0.0
      do k=1,nx
        ! PDE  a1(x,y)*uxx(x,y) term 
	dxx = dxx + cxx(i,k)*ug(k,j)
        ! PDE  d1(x,y)*ux(x,y) term 
	dx = dx + cx(i,k)*ug(k,j)
      end do
      dyy = 0.0
      dy  = 0.0
      do k=1,ny
        ! PDE c1(x,y)*uyy(x,y) term 
	dyy = dyy + cyy(j,k)*ug(i,k)
        ! PDE e1(x,y)*uy(x,y) term 
	dy = dy + cy(j,k)*ug(i,k)
      end do
      ! PDE b1(x,y)*uxy(x,y) term 
      dxy = 0.0
      do jj=1,ny ! walk jj in y, computing uxx 
	tcxy(jj) = 0.0 ! numeric  ux part of uxy 
        do k=1,nx
	  tcxy(jj) = tcxy(jj) + cx(i,k)*ug(k,jj)
	end do
      end do
      do k=1,ny  ! uy part of uxy 
        dxy = dxy + cy(j,k)*tcxy(k)  ! PDE uxy term 
      end do
      ! check 
      if(ifcheck>5) then
        print "('numerical dxx=',f10.3,', dxy=',f10.3,', dyy=',f10.3)", &
              dxx, dxy, dyy
        print "('analytic  dxx=',f10.3,', dxy=',f10.3,', dyy=',f10.3)", &
	        6.0*x+6.0*y, 6.0*x+8.0*y+5.0, 12.0*y+8.0*x
        print "('numerical dx=',f10.3,', dy=',f10.3,', u='f10.3)", &
                dx, dy, ug(i,j)
        print "('analytic  dx=',f10.3,', dy=',f10.3,',u=',f10.3)", &
               3.0*x*x+6.0*x*y+4.0*y*y+5.0*y+6.0, &
               6.0*y*y+3.0*x*x+8.0*x*y+5.0*x+7.0, &
               u(x,y)
      end if
      err = a1(xg(i),yg(j))*dxx + &
	    b1(xg(i),yg(j))*dxy + &
            c1(xg(i),yg(j))*dyy + &
            d1(xg(i),yg(j))*dx  + &
            e1(xg(i),yg(j))*dy  + &
            f1(xg(i),yg(j))*ug(i,j) - &
	     c(xg(i),yg(j))             ! RHS 
      err = abs(err)
      sumerr = sumerr + err
      sumsqerr = sumsqerr + err*err
      maxerr = max(err,maxerr)
      cnt = cnt + 1
    end do ! j 
  end do ! i 
  rmserr = sqrt(sumsqerr/cnt)
  avgerr = sumerr/cnt
  print *, "check_soln values against PDE "
  print "('maxerr=',e15.6,', rmserr=',e15.6,', avgerr=',e15.6)", &
        maxerr,rmserr,avgerr
end subroutine check_soln  

program pde_abc_eq
  use global
  implicit none
  integer :: i, j, ii, jj
  double precision :: x, y

  print *, 'pde_abc_eq.f90 running'
  print *, 'solve a1(x,y) uxx(x,y) + b1(x,y) uxy(x,y) + c1(x,y) uyy(x,y) +'
  print *, '     d1(x,y) ux(x,y)  + e1(x,y) uy(x,y)  + f1(x,y) u(x,y) = c(x,y)'
  print *, 'high order discretization, three area solution with one method'
  print *, 'boundary conditions computed using u(x)'
  print *, 'analytic solution may be given by u(x) ifcheck>0'
  print *, ' '
  print *, 'xmin=',xmin,', xmax =',xmax
  print *, 'ymin=',ymin,', ymax =',ymax
  print *, 'nx=',nx,', ny =',ny,', nd=',nd

  hx = (xmax-xmin)/(nx-1) ! does not have to be uniform spacing
                          ! deriv would become nuderiv
  hx2 = hx*hx
  ! print *, 'x grid and analytic solution, at ymin'
  do i=1,nx 
    xg(i) = xmin+(i-1)*hx
    ! print *, 'i=',i,', u(',xg(i),',ymin)=',u(xg(i),ymin)
  end do
  ! print *, ' '

  hy = (ymax-ymin)/(ny-1) ! does not have to be uniform spacing
  hy2 = hy*hy
  hxy = hx*hy
  ! print *, 'y grid and analytic solution, at xmin'
  do ii=1,ny 
    yg(ii) = ymin+(ii-1)*hy
    ! print *, 'ii=',ii,', u(xmin,',yg(ii),')=',u(xmin,yg(ii))
  end do
  ! print *, ' '
  cs = nxy
  md = nd/2+1
  print *, "hx=",hx,", hy=", hy

  call clear()  
  call init_matrix()
  print *, "matrix initialized"
  if(ifcheck>0) call print_mat()

  call simeq((nx-2)*(ny-2), (nx-2)*(ny-2), ut, uy, us)
  call printus()
  if(ifcheck>0) then
    call printexact()
    call printdiff()
    call check()
    avgerr = error/((nx-1)*(ny-1))
    print "('error=',e15.8,' avg=',e15.8,' max=',e15.8)", &
             error, avgerr, maxerr
  else
    !  derivative computation on solution and check against PDE
    print *, "check solution against PDE"
    call check_soln()
  end if
  print *, "check solution against PDE" ! here also for test
  call check_soln()
end program pde_abc_eq