!
!   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.
!

! -----------------------------------------------------------------------
program main
#include <petsc/finclude/petscksp.h>
  use petscksp
  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, nel, ll, j1, i1, ndf, f4
  PetscInt f2, f9, f6, one
  PetscInt :: idx(4)
  PetscBool flg, out_matlab
  PetscMPIInt size, rank
  PetscScalar::ss(4, 4), val
  PetscReal::shp(3, 9), sg(3, 9)
  PetscReal::thk, a1, a2
  PetscReal::theta, eps, h, x, y, xsj
  PetscReal::coord(2, 4), dd(2, 2), ev(3), blb(2)

  common/ex54_theta/theta
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                 Beginning of program
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  PetscCallA(PetscInitialize(ierr))
  PetscCallMPIA(MPI_Comm_size(PETSC_COMM_WORLD, size, ierr))
  PetscCallMPIA(MPI_Comm_rank(PETSC_COMM_WORLD, rank, ierr))
  one = 1
! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!                 set parameters
!     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  f4 = 4
  f2 = 2
  f9 = 9
  f6 = 6
  ne = 9
  PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-ne', ne, flg, ierr))
  h = 2.0/real(ne)
  M = (ne + 1)*(ne + 1)
  theta = 90.0
!     theta is input in degrees
  PetscCallA(PetscOptionsGetReal(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-theta', theta, flg, ierr))
  theta = theta/57.2957795
  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(1) = 0.0
    blb(2) = 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, f9, PETSC_NULL_INTEGER_ARRAY, f6, 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").
  thk = 1.0              ! thickness
  nel = 4                   ! nodes per element (quad)
  ndf = 1
  call int2d(f2, 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, f2)
        xsj = xsj*sg(3, ll)*thk
        call thfx2d(ev, coord, shp, dd, f2, f2, f4, 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(1) = geq; idx(2) = geq + 1; idx(3) = geq + (ne + 1) + 1
      idx(4) = geq + (ne + 1)
      if (qj > 0) then
        PetscCallA(MatSetValues(Amat, f4, idx, f4, idx, reshape(ss, [f4*f4]), 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, f4, idx, f4, idx, reshape(ss, [f4*f4]), 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, one, [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))

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

  subroutine thfx2d(ev, xl, shp, dd, ndm, ndf, nel, dir)

    PetscInt ndm, ndf, nel, i
    PetscReal ev(2), xl(ndm, nel), shp(3, *), dir
    PetscReal xx, yy, psi, cs, sn, c2, s2, dd(2, 2)

    xx = 0.0
    yy = 0.0
    do i = 1, nel
      xx = xx + shp(3, i)*xl(1, i)
      yy = yy + shp(3, i)*xl(2, i)
    end do
    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

!     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!     shp2dquad - shape functions - compute derivatives w/r natural coords.
!     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
  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 ndm
    PetscReal xo, xs, xt, yo, ys, yt, xsm, xsp, xtm
    PetscReal xtp, ysm, ysp, ytm, ytp
    PetscReal s, t, xsj, xsj1, sh, th, sp, tp, sm
    PetscReal tm, xl(ndm, 4), shp(3, 4)

!     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

!     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!     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, lr(9), lz(9)
    PetscReal g, third, sg(3, *)
    data lr/-1, 1, 1, -1, 0, 1, 0, -1, 0/, lz/-1, -1, 1, 1, -1, 0, 1, 0, 0/
    data third/0.3333333333333333/

!     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 function ex54_psi(x, y)
    PetscReal x, y, theta
    common/ex54_theta/theta
    ex54_psi = theta
    if (theta < 0.) then     ! circular
      if (y == 0) then
        ex54_psi = 2.0*atan(1.0)
      else
        ex54_psi = atan(-x/y)
      end if
    end if
  end
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*/
