!
!   Description: Solve Ax=b.  A comes from an anisotropic 2D thermal problem with Q1 FEM on domain (-1,1)^2.
!       Material conductivity given by tensor:
!
!       D = | 1 0       |
!           | 0 epsilon |
!
!    rotated by angle 'theta' (-theta <90> in degrees) with anisotropic parameter 'epsilon' (-epsilon <0.0>).
!    Blob right-hand side centered at C (-blob_center C(1),C(2) <0,0>)
!    Dirichlet BCs on y=-1 face.
!
!    -out_matlab will generate binary files for A,x,b and a ex54f.m file that reads them and plots them in matlab.
!
!    User can change anisotropic shape with function ex54_psi().  Negative theta will switch to a circular anisotropy.
!

! -----------------------------------------------------------------------
#include <petsc/finclude/petscksp.h>
module ex54fmodule
  use petscksp

  implicit none
  PetscReal theta

contains
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!     thfx2d - compute material tensor
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!     Compute thermal gradient and flux

  pure subroutine thfx2d(ev, xl, shp, dd, ndm, nel, dir)
    PetscInt, intent(in) :: nel, ndm
    PetscReal, intent(out) :: dd(2, 2)
    PetscReal, intent(in) :: ev(2), xl(ndm, nel), shp(3, *)
    procedure(ex54_psi) :: dir
    PetscReal xx, yy, psi, cs, sn, c2, s2

    xx = sum(shp(3, 1:nel)*xl(1, 1:nel))
    yy = sum(shp(3, 1:nel)*xl(2, 1:nel))
    psi = dir(xx, yy)
!   Compute thermal flux
    cs = cos(psi)
    sn = sin(psi)
    c2 = cs*cs
    s2 = sn*sn
    cs = cs*sn

    dd(1, 1) = c2*ev(1) + s2*ev(2)
    dd(2, 2) = s2*ev(1) + c2*ev(2)
    dd(1, 2) = cs*(ev(1) - ev(2))
    dd(2, 1) = dd(1, 2)

!   flux(1) = -dd(1,1)*gradt(1) - dd(1,2)*gradt(2)
!   flux(2) = -dd(2,1)*gradt(1) - dd(2,2)*gradt(2)
  end subroutine thfx2d

!     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!     shp2dquad - shape functions - compute derivatives w/r natural coords.
!     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  pure subroutine shp2dquad(s, t, xl, shp, xsj, ndm)
!-----[--.----+----.----+----.-----------------------------------------]
!      Purpose: Shape function routine for 4-node isoparametric quads
!
!      Inputs:
!         s,t       - Natural coordinates of point
!         xl(ndm,*) - Nodal coordinates for element
!         ndm       - Spatial dimension of mesh

!      Outputs:
!         shp(3,*)  - Shape functions and derivatives at point
!                     shp(1,i) = dN_i/dx  or dN_i/dxi_1
!                     shp(2,i) = dN_i/dy  or dN_i/dxi_2
!                     shp(3,i) = N_i
!         xsj       - Jacobian determinant at point
!-----[--.----+----.----+----.-----------------------------------------]
    PetscInt, intent(in) :: ndm
    PetscReal, intent(out) :: shp(3, 4), xsj
    PetscReal, intent(in) :: xl(ndm, 4), s, t
    PetscReal xo, xs, xt, yo, ys, yt, xsm, xsp, xtm
    PetscReal xtp, ysm, ysp, ytm, ytp
    PetscReal xsj1, sh, th, sp, tp, sm, tm

!   Set up interpolations

    sh = 0.5*s
    th = 0.5*t
    sp = 0.5 + sh
    tp = 0.5 + th
    sm = 0.5 - sh
    tm = 0.5 - th
    shp(3, 1) = sm*tm
    shp(3, 2) = sp*tm
    shp(3, 3) = sp*tp
    shp(3, 4) = sm*tp

!   Set up natural coordinate functions (times 4)

    xo = xl(1, 1) - xl(1, 2) + xl(1, 3) - xl(1, 4)
    xs = -xl(1, 1) + xl(1, 2) + xl(1, 3) - xl(1, 4) + xo*t
    xt = -xl(1, 1) - xl(1, 2) + xl(1, 3) + xl(1, 4) + xo*s
    yo = xl(2, 1) - xl(2, 2) + xl(2, 3) - xl(2, 4)
    ys = -xl(2, 1) + xl(2, 2) + xl(2, 3) - xl(2, 4) + yo*t
    yt = -xl(2, 1) - xl(2, 2) + xl(2, 3) + xl(2, 4) + yo*s

!   Compute jacobian (times 16)

    xsj1 = xs*yt - xt*ys

!   Divide jacobian by 16 (multiply by .0625)

    xsj = 0.0625*xsj1
    if (xsj1 == 0.0) then
      xsj1 = 1.0
    else
      xsj1 = 1.0/xsj1
    end if

!   Divide functions by jacobian

    xs = (xs + xs)*xsj1
    xt = (xt + xt)*xsj1
    ys = (ys + ys)*xsj1
    yt = (yt + yt)*xsj1

!   Multiply by interpolations

    ytm = yt*tm
    ysm = ys*sm
    ytp = yt*tp
    ysp = ys*sp
    xtm = xt*tm
    xsm = xs*sm
    xtp = xt*tp
    xsp = xs*sp

!   Compute shape functions

    shp(1, 1) = -ytm + ysm
    shp(1, 2) = ytm + ysp
    shp(1, 3) = ytp - ysp
    shp(1, 4) = -ytp - ysm
    shp(2, 1) = xtm - xsm
    shp(2, 2) = -xtm - xsp
    shp(2, 3) = -xtp + xsp
    shp(2, 4) = xtp + xsm

  end subroutine shp2dquad

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
! int2d
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  subroutine int2d(l, sg)
!-----[--.----+----.----+----.-----------------------------------------]
!     Purpose: Form Gauss points and weights for two dimensions

!     Inputs:
!     l       - Number of points/direction

!     Outputs:
!     sg(3,*) - Array of points and weights
!-----[--.----+----.----+----.-----------------------------------------]
    PetscInt l, i
    PetscReal g, sg(3, *)
    PetscInt, parameter, dimension(9) :: &
      lr = [-1, 1, 1, -1, 0, 1, 0, -1, 0], &
      lz = [-1, -1, 1, 1, -1, 0, 1, 0, 0]
    PetscReal, parameter :: third = 1.0_PETSC_REAL_KIND/3.0_PETSC_REAL_KIND

!   2x2 integration
    g = sqrt(third)
    do i = 1, 4
      sg(1, i) = g*lr(i)
      sg(2, i) = g*lz(i)
      sg(3, i) = 1.0
    end do

  end

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
! ex54_psi - anisotropic material direction
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  PetscReal pure function ex54_psi(x, y)
    PetscReal, intent(in) :: x, y
    ex54_psi = theta
    if (theta < 0.) then     ! circular
      if (y == 0.0) then
        ex54_psi = 2.0*atan(1.0)
      else
        ex54_psi = atan(-x/y)
      end if
    end if
  end
end module ex54fmodule

program main
  use petscksp
  use ex54fmodule
  implicit none

  Vec xvec, bvec, uvec
  Mat Amat
  KSP ksp
  PetscErrorCode ierr
  PetscViewer viewer
  PetscInt qj, qi, ne, M, Istart, Iend, geq
  PetscInt ki, kj, ll, j1, i1
  PetscInt, parameter :: ndf = 1
  PetscInt, parameter :: nel = 4 ! nodes per element (quad)
  PetscInt :: idx(4)
  PetscBool flg, out_matlab
  PetscMPIInt size, rank
  PetscScalar :: ss(4, 4), val
  PetscReal :: shp(3, 9), sg(3, 9)
  PetscReal :: thk = 1.0              ! thickness
  PetscReal :: eps, h, x, y, xsj, a1, a2
  PetscReal :: coord(2, 4), dd(2, 2), ev(3), blb(2)
  PetscReal, parameter :: rad2deg = 57.2957795_PETSC_REAL_KIND

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                 Beginning of program
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  PetscCallA(PetscInitialize(ierr))
  PetscCallMPIA(MPI_Comm_size(PETSC_COMM_WORLD, size, ierr))
  PetscCallMPIA(MPI_Comm_rank(PETSC_COMM_WORLD, rank, ierr))
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                 set parameters
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  ne = 9
  PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-ne', ne, flg, ierr))
  h = 2.0/real(ne)
  M = (ne + 1)**2
  theta = 90.0
! theta is input in degrees
  PetscCallA(PetscOptionsGetReal(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-theta', theta, flg, ierr))
  theta = theta/rad2deg
  eps = 1.0
  PetscCallA(PetscOptionsGetReal(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-epsilon', eps, flg, ierr))
  ki = 2
  PetscCallA(PetscOptionsGetRealArray(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-blob_center', blb, ki, flg, ierr))
  if (.not. flg) then
    blb = 0.0
  else if (ki /= 2) then
    print *, 'error: ', ki, ' arguments read for -blob_center.  Needs to be two.'
  end if
  PetscCallA(PetscOptionsGetBool(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-out_matlab', out_matlab, flg, ierr))
  if (.not. flg) out_matlab = PETSC_FALSE

  ev(1) = 1.0
  ev(2) = eps*ev(1)
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!    Compute the matrix and right-hand-side vector that define
!    the linear system, Ax = b.
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
! Create matrix.  When using MatCreate(), the matrix format can
! be specified at runtime.
  PetscCallA(MatCreate(PETSC_COMM_WORLD, Amat, ierr))
  PetscCallA(MatSetSizes(Amat, PETSC_DECIDE, PETSC_DECIDE, M, M, ierr))
  PetscCallA(MatSetType(Amat, MATAIJ, ierr))
  PetscCallA(MatSetOption(Amat, MAT_SPD, PETSC_TRUE, ierr))
  PetscCallA(MatSetOption(Amat, MAT_SPD_ETERNAL, PETSC_TRUE, ierr))
  if (size == 1) then
    PetscCallA(MatSetType(Amat, MATAIJ, ierr))
  else
    PetscCallA(MatSetType(Amat, MATMPIAIJ, ierr))
  end if
  PetscCallA(MatMPIAIJSetPreallocation(Amat, 9_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, 6_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, ierr))
  PetscCallA(MatSetFromOptions(Amat, ierr))
  PetscCallA(MatSetUp(Amat, ierr))
  PetscCallA(MatGetOwnershipRange(Amat, Istart, Iend, ierr))
!  Create vectors.  Note that we form 1 vector from scratch and
!  then duplicate as needed.
  PetscCallA(MatCreateVecs(Amat, PETSC_NULL_VEC, xvec, ierr))
  PetscCallA(VecSetFromOptions(xvec, ierr))
  PetscCallA(VecDuplicate(xvec, bvec, ierr))
  PetscCallA(VecDuplicate(xvec, uvec, ierr))
! Assemble matrix.
!  - Note that MatSetValues() uses 0-based row and column numbers
!    in Fortran as well as in C (as set here in the array "col").
  call int2d(2_PETSC_INT_KIND, sg)
  do geq = Istart, Iend - 1, 1
    qj = geq/(ne + 1); qi = mod(geq, (ne + 1))
    x = h*qi - 1.0; y = h*qj - 1.0 ! lower left corner (-1,-1)
    if (qi < ne .and. qj < ne) then
      coord(1, 1) = x; coord(2, 1) = y
      coord(1, 2) = x + h; coord(2, 2) = y
      coord(1, 3) = x + h; coord(2, 3) = y + h
      coord(1, 4) = x; coord(2, 4) = y + h
! form stiff
      ss = 0.0
      do ll = 1, 4
        call shp2dquad(sg(1, ll), sg(2, ll), coord, shp, xsj, 2_PETSC_INT_KIND)
        xsj = xsj*sg(3, ll)*thk
        call thfx2d(ev, coord, shp, dd, 2_PETSC_INT_KIND, 4_PETSC_INT_KIND, ex54_psi)
        j1 = 1
        do kj = 1, nel
          a1 = (dd(1, 1)*shp(1, kj) + dd(1, 2)*shp(2, kj))*xsj
          a2 = (dd(2, 1)*shp(1, kj) + dd(2, 2)*shp(2, kj))*xsj
!     Compute residual
!         p(j1) = p(j1) - a1*gradt(1) - a2*gradt(2)
!     Compute tangent
          i1 = 1
          do ki = 1, nel
            ss(i1, j1) = ss(i1, j1) + a1*shp(1, ki) + a2*shp(2, ki)
            i1 = i1 + ndf
          end do
          j1 = j1 + ndf
        end do
      end do

      idx = [geq, geq + 1, geq + (ne + 1) + 1, geq + (ne + 1)]
      if (qj > 0) then
        PetscCallA(MatSetValues(Amat, 4_PETSC_INT_KIND, idx, 4_PETSC_INT_KIND, idx, reshape(ss, [4_PETSC_INT_KIND**2]), ADD_VALUES, ierr))
      else                !     a BC
        do ki = 1, 4, 1
          do kj = 1, 4, 1
            if (ki < 3 .or. kj < 3) then
              if (ki == kj) then
                ss(ki, kj) = .1*ss(ki, kj)
              else
                ss(ki, kj) = 0.0
              end if
            end if
          end do
        end do
        PetscCallA(MatSetValues(Amat, 4_PETSC_INT_KIND, idx, 4_PETSC_INT_KIND, idx, reshape(ss, [4_PETSC_INT_KIND**2]), ADD_VALUES, ierr))
      end if               ! BC
    end if                  ! add element
    if (qj > 0) then      ! set rhs
      val = h*h*exp(-100*((x + h/2) - blb(1))**2)*exp(-100*((y + h/2) - blb(2))**2)
      PetscCallA(VecSetValues(bvec, 1_PETSC_INT_KIND, [geq], [val], INSERT_VALUES, ierr))
    end if
  end do
  PetscCallA(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY, ierr))
  PetscCallA(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY, ierr))
  PetscCallA(VecAssemblyBegin(bvec, ierr))
  PetscCallA(VecAssemblyEnd(bvec, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!          Create the linear solver and set various options
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

!  Create linear solver context

  PetscCallA(KSPCreate(PETSC_COMM_WORLD, ksp, ierr))

!  Set operators. Here the matrix that defines the linear system
!  also serves as the matrix from which the preconditioner is constructed.

  PetscCallA(KSPSetOperators(ksp, Amat, Amat, ierr))

!  Set runtime options, e.g.,
!      -ksp_type <type> -pc_type <type> -ksp_monitor -ksp_rtol <rtol>
!  These options will override those specified above as long as
!  KSPSetFromOptions() is called _after_ any other customization
!  routines.

  PetscCallA(KSPSetFromOptions(ksp, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                      Solve the linear system
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

  PetscCallA(KSPSolve(ksp, bvec, xvec, ierr))

! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                      output
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  if (out_matlab) then
    PetscCallA(PetscViewerBinaryOpen(PETSC_COMM_WORLD, 'Amat', FILE_MODE_WRITE, viewer, ierr))
    PetscCallA(MatView(Amat, viewer, ierr))
    PetscCallA(PetscViewerDestroy(viewer, ierr))

    PetscCallA(PetscViewerBinaryOpen(PETSC_COMM_WORLD, 'Bvec', FILE_MODE_WRITE, viewer, ierr))
    PetscCallA(VecView(bvec, viewer, ierr))
    PetscCallA(PetscViewerDestroy(viewer, ierr))

    PetscCallA(PetscViewerBinaryOpen(PETSC_COMM_WORLD, 'Xvec', FILE_MODE_WRITE, viewer, ierr))
    PetscCallA(VecView(xvec, viewer, ierr))
    PetscCallA(PetscViewerDestroy(viewer, ierr))

    PetscCallA(MatMult(Amat, xvec, uvec, ierr))
    val = -1.0
    PetscCallA(VecAXPY(uvec, val, bvec, ierr))
    PetscCallA(PetscViewerBinaryOpen(PETSC_COMM_WORLD, 'Rvec', FILE_MODE_WRITE, viewer, ierr))
    PetscCallA(VecView(uvec, viewer, ierr))
    PetscCallA(PetscViewerDestroy(viewer, ierr))

    if (rank == 0) then
      open (1, file='ex54f.m', FORM='formatted')
      write (1, *) 'A = PetscBinaryRead(''Amat'');'
      write (1, *) '[m n] = size(A);'
      write (1, *) 'mm = sqrt(m);'
      write (1, *) 'b = PetscBinaryRead(''Bvec'');'
      write (1, *) 'x = PetscBinaryRead(''Xvec'');'
      write (1, *) 'r = PetscBinaryRead(''Rvec'');'
      write (1, *) 'bb = reshape(b,mm,mm);'
      write (1, *) 'xx = reshape(x,mm,mm);'
      write (1, *) 'rr = reshape(r,mm,mm)'
!     write (1,*) 'imagesc(bb')'
!     write (1,*) 'title('RHS'),'
      write (1, *) 'figure,'
      write (1, *) 'imagesc(xx'')'
      write (1, 2002) eps, theta*57.2957795
      write (1, *) 'figure,'
      write (1, *) 'imagesc(rr'')'
      write (1, *) 'title(''Residual''),'
      close (1)
    end if
  end if
2002 format('title(''Solution: esp='',d9.3,'', theta='',g8.3,''),')
!  Free work space.  All PETSc objects should be destroyed when they
!  are no longer needed.

  PetscCallA(VecDestroy(xvec, ierr))
  PetscCallA(VecDestroy(bvec, ierr))
  PetscCallA(VecDestroy(uvec, ierr))
  PetscCallA(MatDestroy(Amat, ierr))
  PetscCallA(KSPDestroy(ksp, ierr))
  PetscCallA(PetscFinalize(ierr))
end

!
!/*TEST
!
!  testset:
!   nsize: 4
!   args: -ne 39 -theta 30.0 -epsilon 1.e-1 -blob_center 0.,0. -ksp_type cg -pc_type gamg -pc_gamg_type agg -ksp_rtol 1e-4 -pc_gamg_aggressive_square_graph false -ksp_norm_type unpreconditioned
!   requires: !single
!   test:
!      suffix: misk
!      args: -pc_gamg_mat_coarsen_type misk -pc_gamg_aggressive_coarsening 0 -ksp_monitor_short
!   test:
!      suffix: mis
!      args: -pc_gamg_mat_coarsen_type mis -ksp_monitor_short
!   test:
!      suffix: hem
!      args: -pc_gamg_mat_coarsen_type hem -ksp_converged_reason
!      filter: sed -e "s/Linear solve converged due to CONVERGED_RTOL iterations 1[2-3]/Linear solve converged due to CONVERGED_RTOL iterations 11/g"
!
!TEST*/
