!==========================================================!
! Differential operator routines !
!==========================================================!
! SUBROUTINE div_3d(ax, ay, az, div_a) !
! SUBROUTINE curl_3d(ax, ay, az, curl_x, curl_y, curl_z) !
! SUBROUTINE grad_3d(a, grad_x, grad_y, grad_z) !
! SUBROUTINE init_obj_funct() !
! FUNCTION obj_funct_fff(b2, omega2) !
! SUBROUTINE calc_lin_fff() !
!==========================================================!
MODULE fff_operators
  IMPLICIT NONE

  CONTAINS

  !--------------------------------------------------------------!
  ! SUBROUTINE div_3d !
  !--------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates the divergence for a 3-d vector !
  ! INPUT: !
  ! ax,ay,az = vector components of a; dimension (nx,ny,nz) !
  ! OUTPUT: !
  ! div_a = divergence of a; dimension (nx,ny,nz) !
  ! REQUIRES: !
  ! x,y,z,nx,ny,nz,qsphere: from vars !
  ! HISTORY: !
  ! 22-may-2006: jmm, jimm@ssl.berkeley.edu !
  ! 07-feb-2007: jmm, try to allocate minimal memory !
  ! Feb-2008: djb, modularized, based on earlier versions !
  !--------------------------------------------------------------!
  SUBROUTINE div_3d(ax, ay, az, div_a)
    USE params, ONLY: i4,r8
    USE vars, ONLY: x,y,z,nx,ny,nz,qsphere
    USE fff_derivs
    IMPLICIT NONE

    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: ax, ay, az
    REAL(KIND=r8), DIMENSION(SIZE(x,1),SIZE(y,1),SIZE(z,1)), &
     INTENT(OUT) :: div_a
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: sinth, cotth
    REAL(KIND=r8), DIMENSION(:,:,:), ALLOCATABLE :: f1

    IF(qsphere .EQV. .TRUE.) THEN ! spherical coordinates
      ALLOCATE(f1(nx,ny,nz))
      ALLOCATE(sinth(ny))
      ALLOCATE(cotth(ny))
      sinth = SIN(y)
      cotth = COS(y)/sinth
      f1 = deriv_dx(ax)
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        div_a(i,j,k) = f1(i,j,k) + 2.0_r8*ax(i,j,k)/x(i)
      END FORALL
      f1 = deriv_dy(ay)
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        div_a(i,j,k) = div_a(i,j,k) &
                       + (f1(i,j,k) + cotth(j)*ay(i,j,k))/x(i)
      END FORALL
      f1 = deriv_dz(az)
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        div_a(i,j,k) = div_a(i,j,k) + f1(i,j,k)/(x(i)*sinth(j))
      END FORALL
      DEALLOCATE(f1, sinth, cotth)
    ELSE ! Cartesian coordinates
      div_a = deriv_dx(ax) + deriv_dy(ay) + deriv_dz(az)
    ENDIF
  END SUBROUTINE div_3d

  !--------------------------------------------------------------!
  ! SUBROUTINE curl_3d !
  !--------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates the curl for a 3-d vector !
  ! INPUT: !
  ! ax,ay,az = vector components of a; dimension (nx,ny,nz) !
  ! OUTPUT: !
  ! curl_x,curl_y,curl_z = vector components of the curl of a; !
  ! dimension (nx,ny,nz) !
  ! REQUIRES: !
  ! x,y,z,nx,ny,nz,qsphere: from vars !
  ! HISTORY: !
  ! 18-may-2006: jmm, jimm@ssl.berkeley.edu !
  ! 07-feb-2007: jmm, trying to save some memory !
  ! Feb-2008: djb, modularized, based on earlier versions !
  !--------------------------------------------------------------!
  SUBROUTINE curl_3d(ax, ay, az, curl_x, curl_y, curl_z)
    USE params, ONLY: i4,r8
    USE vars, ONLY: x,y,z,nx,ny,nz,qsphere
    USE fff_derivs
    IMPLICIT NONE

    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: ax, ay, az
    REAL(KIND=r8), DIMENSION(SIZE(x,1),SIZE(y,1),SIZE(z,1)), &
     INTENT(OUT) :: curl_x, curl_y, curl_z
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: sinth
    REAL(KIND=r8), DIMENSION(:,:,:), ALLOCATABLE :: f1

    IF(qsphere .EQV. .TRUE.) THEN
      ALLOCATE(f1(nx,ny,nz))
      ALLOCATE(sinth(ny))
      sinth = sin(y)
      ! r-component
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        f1(i,j,k) = sinth(j)*az(i,j,k)
      END FORALL
      curl_x = deriv_dy(f1)
      f1 = deriv_dz(ay)
      curl_x = curl_x - f1
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        curl_x(i,j,k) = curl_x(i,j,k)/(x(i)*sinth(j))
      END FORALL
      ! theta-component
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        f1(i,j,k) = x(i)*az(i,j,k)
      END FORALL
      curl_y = deriv_dx(f1)
      f1 = deriv_dz(ax)
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        curl_y(i,j,k) = (f1(i,j,k)/sinth(j) - curl_y(i,j,k))/x(i)
      END FORALL
      ! phi-component
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        f1(i,j,k) = x(i)*ay(i,j,k)
      END FORALL
      curl_z = deriv_dx(f1)
      f1 = deriv_dy(ax)
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        curl_z(i,j,k) = (curl_z(i,j,k) - f1(i,j,k))/x(i)
      END FORALL
      DEALLOCATE(f1, sinth)
    ELSE
      curl_x = deriv_dy(az) - deriv_dz(ay)
      curl_y = deriv_dz(ax) - deriv_dx(az)
      curl_z = deriv_dx(ay) - deriv_dy(ax)
    END IF
  END SUBROUTINE curl_3d

  !--------------------------------------------------------------!
  ! SUBROUTINE grad_3d !
  !--------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates the gradient for a 3-d array !
  ! INPUT: !
  ! a = 3-d array; dimension (nx,ny,nz) !
  ! OUTPUT: !
  ! grad_x, grad_y, grad_z = components of the gradient of a; !
  ! dimension (nx,ny,nz) !
  ! REQUIRES: !
  ! x,y,z,nx,ny,nz,qsphere: from vars !
  ! HISTORY: !
  ! 18-may-2006: jmm, jimm@ssl.berkeley.edu !
  ! Feb-2008: djb, modularized, based on earlier versions !
  !--------------------------------------------------------------!
  SUBROUTINE grad_3d(a, grad_x, grad_y, grad_z)
    USE params, ONLY: i4,r8
    USE vars, ONLY: x,y,z,nx,ny,nz,qsphere
    USE fff_derivs
    IMPLICIT NONE

    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: a
    REAL(KIND=r8), DIMENSION(SIZE(x,1),SIZE(y,1),SIZE(z,1)), &
     INTENT(OUT) :: grad_x, grad_y, grad_z
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: sinth

    grad_x = deriv_dx(a)
    grad_y = deriv_dy(a)
    grad_z = deriv_dz(a)
    IF(qsphere .EQV. .TRUE.) THEN
      ALLOCATE(sinth(ny))
      sinth = SIN(y)
      FORALL(i=1:nx, j=1:ny, k=1:nz)
        grad_y(i,j,k) = grad_y(i,j,k)/x(i)
        grad_z(i,j,k) = grad_z(i,j,k)/(x(i)*sinth(j))
      END FORALL
      DEALLOCATE(sinth)
    END IF
  END SUBROUTINE grad_3d

  !------------------------------------------------------------!
  ! SUBROUTINE init_obj_funct !
  !------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculate grid-based coefficients for the integral in !
  ! obj_funct_fff !
  ! INPUT: !
  ! none !
  ! OUTPUT: !
  ! none; indirect (xint_coeff, yint_coeff, zint_coeff) !
  ! REQUIRES: !
  ! nx,ny,ny,x,y,z,xint_coeff,yint_coeff,zint_coeff,qsphere: !
  ! from vars !
  ! HISTORY: !
  ! 19-May-2006, jmm !
  ! Feb-2008: djb, modularized, based on earlier versions !
  !------------------------------------------------------------!
  SUBROUTINE init_obj_funct
    USE params, ONLY: r8
    USE vars, ONLY: nx, ny, nz, x, y, z, xint_coeff, yint_coeff, &
                    zint_coeff, qsphere
    IMPLICIT NONE

    ! Local variable declarations
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: dx(:), dy(:), dz(:)

    ALLOCATE (dx(nx-1), dy(ny-1), dz(nz-1))
    dx = x(2:nx) - x(1:nx-1)
    dy = y(2:ny) - y(1:ny-1)
    dz = z(2:nz) - z(1:nz-1)
    ! Arrays for integration
    xint_coeff(1) = 0.5_r8*dx(1)
    xint_coeff(nx) = 0.5_r8*dx(nx-1)
    xint_coeff(2:nx-1) = 0.5_r8*(dx(1:nx-2) + dx(2:nx-1))
    yint_coeff(1) = 0.5_r8*dy(1)
    yint_coeff(ny) = 0.5_r8*dy(ny-1)
    yint_coeff(2:ny-1) = 0.5_r8*(dy(1:ny-2) + dy(2:ny-1))
    zint_coeff(1) = 0.5_r8*dz(1)
    zint_coeff(nz) = 0.5_r8*dz(nz-1)
    zint_coeff(2:nz-1) = 0.5_r8*(dz(1:nz-2) + dz(2:nz-1))
    IF(qsphere .EQV. .TRUE.) THEN
      xint_coeff = (x*x)*xint_coeff
      yint_coeff = SIN(y)*yint_coeff
    ENDIF
    DEALLOCATE(dx, dy, dz)
  END SUBROUTINE init_obj_funct

  !-------------------------------------------------------------------!
  ! FUNCTION obj_funct_fff !
  !-------------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates the objective function for the FFF code, see !
  ! Wheatland, et al, 2000, eqn(5). The objective function is the !
  ! integral of b2*Omega2 over the whole volume. !
  ! NOTE: init_obj_funct must be called prior to using this !
  ! function !
  ! INPUT: !
  ! b2: B^2; dimensions (nx,ny,nz) !
  ! omega2: omega^2, where omega = (1/B2)*((curlB X B) - (div_B)B); !
  ! dimensions (nx,ny,nz) !
  ! OUTPUT: !
  ! obj_funct_fff: the integral of B2*Omega2 over the volume; !
  ! dimensions (scalar) !
  ! REQUIRES: !
  ! nx,ny,nz,xint_coeff,yint_coeff,zint_coeff: from vars !
  ! HISTORY: !
  ! 19-May-2006, jmm !
  ! Feb-2008: djb, modularized !
  !-------------------------------------------------------------------!
  FUNCTION obj_funct_fff(b2, omega2)
    USE params, ONLY: i4, r8
    USE vars, ONLY: nx, ny, nz, xint_coeff, yint_coeff, zint_coeff
    IMPLICIT NONE
    !
    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: b2, omega2
    REAL(KIND=r8) :: obj_funct_fff
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k
    !
    obj_funct_fff = 0.0_r8
    DO k=1,nz
      DO j=1,ny
        DO i=1,nx
          obj_funct_fff = obj_funct_fff + b2(i,j,k)*omega2(i,j,k)* &
                          xint_coeff(i)*yint_coeff(j)*zint_coeff(k)
        END DO
      END DO
    END DO
  END FUNCTION obj_funct_fff

  !------------------------------------------------------------!
  ! SUBROUTINE calc_lin_fff !
  !------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculation to determine linear FFF from the field on !
  ! the lower boundary. !
  ! NOTE: Only works for Cartesian coordinates !
  ! INPUT: !
  ! none !
  ! OUTPUT: !
  ! none; indirect (bx, by, bz) !
  ! REQUIRES: !
  ! nx,ny,ny,x,y,z,bx,by,bz,alpha: from vars !
  ! HISTORY: !
  ! 19-May-2006, jmm !
  ! Feb-2008: djb, modularized !
  !------------------------------------------------------------!
  SUBROUTINE calc_lin_fff
    USE params, ONLY: i4, r8, twopi
    USE vars, ONLY: nx, ny, nz, x, y, z, bx, by, bz, alpha
    IMPLICIT NONE

    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k, ip, jp
    REAL(KIND=r8) :: r, bigr, cos_ar, sin_ar, xsum, ysum, zsum
    REAL(KIND=r8) :: g, dgdz, gx, gy, gz, zk, twopiinv
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: dx(:), dy(:)
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: ddx(:), ddy(:)
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: cos_az(:), sin_az(:)

    IF(alpha /= 0.0_r8) THEN
      WRITE(UNIT=6,FMT=*) 'Linear FFF implemented, alpha = ', alpha
    ELSE
      WRITE(UNIT=6,FMT=*) 'Potential Field implemented'
    END IF
    ! Arrays for integration
    twopiinv = 1.0_r8/twopi
    ALLOCATE(dx(nx), dy(ny))
    ALLOCATE(ddx(nx), ddy(ny))
    ALLOCATE(cos_az(nz), sin_az(nz))
    dx(1:nx-1) = x(2:nx) - x(1:nx-1)
    dy(1:ny-1) = y(2:ny) - y(1:ny-1)
    ddx(1) = 0.5_r8*dx(1)
    ddx(nx) = 0.5_r8*dx(nx-1)
    ddx(2:nx-1) = 0.5_r8*(dx(1:nx-2) + dx(2:nx-1))
    ddy(1) = 0.5_r8*dy(1)
    ddy(ny) = 0.5_r8*dy(ny-1)
    ddy(2:ny-1) = 0.5_r8*(dy(1:ny-2) + dy(2:ny-1))
    ! get cos and sin arrays
    cos_az = COS(alpha*z)
    sin_az = SIN(alpha*z)
    ! 5 loops, for each i,j,k, integrate over x and y
    DO k=2,nz
      zk = z(k) - z(1)
      DO j=1,ny
        DO i=1,nx
          xsum = 0.0_r8
          ysum = 0.0_r8
          zsum = 0.0_r8
          DO jp=1,ny
            DO ip=1,nx
              IF((ip /= i) .OR. (jp /= j)) THEN
                bigR = SQRT((x(i) - x(ip))**2 + (y(j) - y(jp))**2)
                r = SQRT(bigR**2 + zk**2)
                cos_ar = COS(alpha*r)
                sin_ar = SIN(alpha*r)
                g = zk*cos_ar/(bigR*r) - cos_az(k)/bigR
                dgdz = cos_ar*(1.0_r8/(bigr*r) - zk**2/(bigr*r**3)) &
                       - alpha*zk**2*sin_ar/(bigr*r**2) &
                       + alpha*sin_az(k)/bigr
                gx = bz(ip,jp,1)*((x(i) - x(ip))*dgdz/bigr &
                     + alpha*g*(y(j) - y(jp))/bigr)
                gy = bz(ip,jp,1)*((y(j) - y(jp))*dgdz/bigr &
                     - alpha*g*(x(i) - x(ip))/bigr)
                gz = bz(ip,jp,1)*(zk*cos_ar/r**3 + alpha*zk*sin_ar/r**2)
!c IF((i == 10) .AND. (j == 4) .AND. (k == 2)) THEN
!c IF((ip == 1) .AND. (ip == 1)) THEN
!c WRITE(UNIT=6,FMT=*) x(i), y(j)
!c WRITE(UNIT=6,FMT=*) x(ip), y(jp)
!c WRITE(UNIT=6,FMT=*) bigr, r, cos_ar, sin_ar
!c WRITE(UNIT=6,FMT=*) g, dgdz
!c WRITE(UNIT=6,FMT=*) bz(ip,jp,1)
!c WRITE(UNIT=6,FMT=*) gx, gy, gz
!c WRITE(UNIT=6,FMT=*) ddx(ip), ddy(jp)
!c END IF
!c END IF
                xsum = xsum + gx*ddx(ip)*ddy(jp)
                ysum = ysum + gy*ddx(ip)*ddy(jp)
                zsum = zsum + gz*ddx(ip)*ddy(jp)
              END IF
            END DO
          END DO
          bx(i,j,k) = xsum*twopiinv
          by(i,j,k) = ysum*twopiinv
          bz(i,j,k) = zsum*twopiinv
!c IF((i == 10) .AND. (j == 4) .AND. (k == 2)) THEN
!c WRITE(UNIT=6,FMT=*) bx(i,j,k), by(i,j,k), bz(i,j,k)
!c STOP
!c END IF
        END DO
      END DO
    END DO

    DEALLOCATE(dx, dy)
    DEALLOCATE(ddx, ddy)
    DEALLOCATE(cos_az, sin_az)
  END SUBROUTINE calc_lin_fff

END MODULE fff_operators
