C$Procedure ZZELLBDS ( Create bounding ellipsoids )

      SUBROUTINE ZZELLBDS ( A, B, HMAX, HMIN, AMAX, BMAX, AMIN, BMIN )
      IMPLICIT NONE

C$ Abstract
C
C     SPICE Private routine intended solely for the support of SPICE
C     routines. Users should not call this routine directly due
C     to the volatile nature of this routine.
C
C     Given an oblate spheroid and upper and lower height bounds
C     relative to that spheroid, determine radii of inner and outer
C     bounding spheroids.
C
C$ Disclaimer
C
C     THIS SOFTWARE AND ANY RELATED MATERIALS WERE CREATED BY THE
C     CALIFORNIA INSTITUTE OF TECHNOLOGY (CALTECH) UNDER A U.S.
C     GOVERNMENT CONTRACT WITH THE NATIONAL AERONAUTICS AND SPACE
C     ADMINISTRATION (NASA). THE SOFTWARE IS TECHNOLOGY AND SOFTWARE
C     PUBLICLY AVAILABLE UNDER U.S. EXPORT LAWS AND IS PROVIDED "AS-IS"
C     TO THE RECIPIENT WITHOUT WARRANTY OF ANY KIND, INCLUDING ANY
C     WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A
C     PARTICULAR USE OR PURPOSE (AS SET FORTH IN UNITED STATES UCC
C     SECTIONS 2312-2313) OR FOR ANY PURPOSE WHATSOEVER, FOR THE
C     SOFTWARE AND RELATED MATERIALS, HOWEVER USED.
C
C     IN NO EVENT SHALL CALTECH, ITS JET PROPULSION LABORATORY, OR NASA
C     BE LIABLE FOR ANY DAMAGES AND/OR COSTS, INCLUDING, BUT NOT
C     LIMITED TO, INCIDENTAL OR CONSEQUENTIAL DAMAGES OF ANY KIND,
C     INCLUDING ECONOMIC DAMAGE OR INJURY TO PROPERTY AND LOST PROFITS,
C     REGARDLESS OF WHETHER CALTECH, JPL, OR NASA BE ADVISED, HAVE
C     REASON TO KNOW, OR, IN FACT, SHALL KNOW OF THE POSSIBILITY.
C
C     RECIPIENT BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF
C     THE SOFTWARE AND ANY RELATED MATERIALS, AND AGREES TO INDEMNIFY
C     CALTECH AND NASA FOR ALL THIRD-PARTY CLAIMS RESULTING FROM THE
C     ACTIONS OF RECIPIENT IN THE USE OF THE SOFTWARE.
C
C$ Required_Reading
C
C     None.
C
C$ Keywords
C
C     DEM
C     DSK
C     ELLIPSOID
C     GEOMETRY
C     TOPOGRAPHY
C
C$ Declarations

      DOUBLE PRECISION      A
      DOUBLE PRECISION      B
      DOUBLE PRECISION      HMAX
      DOUBLE PRECISION      HMIN
      DOUBLE PRECISION      AMAX
      DOUBLE PRECISION      BMAX
      DOUBLE PRECISION      AMIN
      DOUBLE PRECISION      BMIN

C$ Brief_I/O
C
C     Variable  I/O  Description
C     --------  ---  --------------------------------------------------
C     A          I   Input spheroid's semi-major axis length.
C     B          I   Input spheroid's semi-minor axis length.
C     HMAX       I   Maximum height relative to input spheroid.
C     HMIN       I   Minimum height relative to input spheroid.
C     AMAX       O   Outer spheroid's semi-major axis length.
C     BMAX       O   Outer spheroid's semi-minor axis length.
C     AMIN       O   Inner spheroid's semi-major axis length.
C     BMIN       O   Inner spheroid's semi-minor axis length.
C
C$ Detailed_Input
C
C     A          is the semi-major axis length of a reference
C                spheroid. A may have any units; B, HMAX, and
C                HMIN must have the same units.
C
C     B          is the semi-minor axis length of a reference
C                spheroid. B must not exceed A.
C
C     HMAX       is an upper bound for a set of heights relative
C                the reference spheroid defined by A and B. HMAX
C                is a signed quantity.
C
C     HMIN       is a lower bound for a set of heights relative
C                the reference spheroid defined by A and B. HMIN
C                is a signed quantity. HMIN must not exceed HMAX.
C                HMIN must be greater than -B.
C
C$ Detailed_Output
C
C     AMAX,
C     BMAX       are, respectively, semi-major and semi-minor axis
C                lengths for an outer bounding spheroid. The set of
C                points at height HMAX relative to the input spheroid
C                defined by A and B lies on or below the spheroid
C                defined by AMAX and BMAX.
C
C                When HMAX is non-negative
C
C                   AMAX  =  A + HMAX
C                   BMAX  =  B + HMAX*(A/B)
C
C                When HMAX is negative
C
C                   AMAX  =  A + HMAX*(B/A)
C                   BMAX  =  B + HMAX
C     AMIN,
C     BMIN       are, respectively, semi-major and semi-minor axis
C                lengths for an inner bounding spheroid. The set of
C                points at height HMIN relative to the input spheroid
C                defined by A and B lies on or above the spheroid
C                defined by AMIN and BMIN.
C
C                When HMIN is non-positive
C
C                   AMIN  =  A + HMIN
C                   BMIN  =  B + HMIN*(A/B)
C
C                When HMIN is positive
C
C                   AMIN  =  A + HMIN*(B/A)
C                   BMIN  =  B + HMIN
C
C$ Parameters
C
C     None.
C
C$ Exceptions
C
C     1)  If the semi-minor axis length B is non-positive, the
C         error SPICE(NONPOSITIVERADIUS) will be signaled.
C
C     2)  If the semi-major axis length A is less than the
C         semi-minor axis length B, the error SPICE(RADIIOUTOFORDER)
C         will be signaled.
C
C     3)  If HMIN is less than or equal to -B, the error
C         SPICE(LOWERBOUNDTOOLOW) will be signaled.
C
C     4)  If the lower height bound HMIN is greater than the
C         upper height bound HMAX, the error SPICE(BOUNDSOUTOFORDER)
C         will be signaled.
C
C$ Files
C
C     None.
C
C$ Particulars
C
C     Ellipsoidal bounding surfaces may be used for volume elements in
C     planetodetic coordinates. The combination of these surfaces,
C     together with cones of constant latitude and planes of constant
C     longitude, may be used to create a simple set of bounding
C     surfaces for such an element. These surfaces can be used to
C     rapidly determine that a given ray does not intersect the volume
C     element.
C
C     For a surface represented by a digital elevation model (DEM),
C     efficient solution of the ray-surface intercept problem is
C     enhanced by determination of a line segment---also called a
C     "chord"---outside of which no solution is possible. One way of
C     generating a chord is to determine the intersection of the ray
C     with inner or outer bounding surfaces. An inner bounding surface
C     is one which has, for every latitude and longitude covered by the
C     DEM, altitude less than or equal to that of the DEM. An outer
C     bounding surface has, for every latitude and longitude covered by
C     the DEM, altitude greater than or equal to the DEM.
C
C     In order for bounding surfaces to be useful, it must be possible
C     to rapidly compute intersections of rays with these surfaces.
C     The bounding surfaces must also lie "close" to the DEM surface;
C     otherwise chords formed by intersections of rays with the bounding
C     surfaces may be longer than necessary.
C
C     Spheroids are a natural choice for bounding surfaces, since DEMs
C     are typically referenced to a spheroidal surface, and because
C     ray-ellipsoid intersections can be rapidly computed by closed-form
C     algebraic methods.
C
C     It might seem that, given a spheroid with semi-major axis length
C     A and semi-minor axis length B, and given a surface having maximum
C     and minimum spheroid-relative heights HMAX and HMIN respectively,
C     that the level surfaces having heights HMAX and HMIN relative to
C     the spheroid would be candidates for the inner and outer bounding
C     spheroids. The outer spheroid would have radii A+HMAX and
C     B+HMAX, while the inner spheroid would have radii A+HMIN and
C     B+HMIN. However, it can be shown (by numerical comparison, for
C     example), that these spheroids do not bound the level surfaces
C     at heights HMAX and HMIN, except in certain special cases such as
C     that of a spherical reference surface.
C
C     This routine generates semi-axis lengths of inner and outer
C     bounding spheroids that are valid for all eccentricities of the
C     reference spheroid (as long as A >= B) and, for reference
C     spheroids of low eccentricity, lie close to the level surfaces at
C     heights HMIN and HMAX. A discussion of the method follows.
C
C
C     Explanation
C     ===========
C
C     Since we're working with spheroids, we can reduce the problem
C     to a two-dimensional one. We'll compute inner and outer bounding
C     ellipses for an ellipse with positive semi-axis lengths A and B.
C     Revolving the bounding ellipses about the Z axis creates bounding
C     ellipsoids.
C
C        Consider an ellipse E with vertical semi-minor axis length
C        B and horizontal semi-major axis A. For a point (x,y) on E,
C        let N be the normal vector
C
C                  x      y
C           N = ( ---- , ---- )
C                   2      2
C                  A      B
C
C        Let LAMBDA be a constant, and Let E' be the curve
C
C           { (x,y) + LAMBDA*N }
C
C        Then E' is an ellipse (since it's produced by a linear
C        transformation of E) with semi-axis lengths
C
C           (A + LAMBDA/A), (B + LAMBDA/B)
C
C
C        For any point (x,y) on E, the height of E' above that point is
C
C           LAMBDA*||N||
C
C        The square of this height HSQ is
C
C
C                       2
C           HSQ = LAMBDA  * <N, N>
C
C
C                               2      2
C                       2      x      y
C               = LAMBDA  * ( ---- + ---- )
C                               4      4
C                              A      B
C
C                               2     2       2  2
C                       2      x     B ( 1 - x /A )
C               = LAMBDA  * ( ---- + -------------- )
C                               4           4
C                              A           B
C
C                               2
C                       2      x     1        1        1
C               = LAMBDA  * ( ---- (----  -  ----) +  ---- )
C                               2     2        2        2
C                              A     A        B        B
C
C
C        If A = B, this expression is constant.
C
C        If A > B, then the term
C
C           1/A**2 - 1/B**2
C
C        is negative, so HSQ is a decreasing function of x for
C
C           0 <= x <= A
C
C        This implies that for x in the above range, E' is closest
C        to E when x = A. We'll use this fact to generate bounding
C        ellipsoids with the following properties:
C
C           If HMAX >= 0    E' will have height HMAX at x = A and
C                           height >= HMAX if 0 <= x < A.
C
C           If HMAX <  0    E' will have height HMAX at x = 0 and
C                           height >= HMAX if 0 < x <= A.
C
C           If HMIN <= 0    E' will have height HMIN at x = A and
C                           height <= HMIN if 0 <= x < A.
C
C           If HMIN >  0    E' will have height HMIN at x = 0 and
C                           height <= HMIN if 0 < x <= A.
C
C
C
C     Application to prolate spheroids
C     ================================
C
C     For a spheroid having semi-axes A, B, C, for which
C
C        A = B < C
C
C     This routine can be applied to the semi-axis lengths
C     A' and B', where
C
C        A'  =  C
C        B'  =  A
C
C     The outer bounding surface is a prolate spheroid with semi-axis
C     lengths
C
C        BMAX, BMAX, AMAX
C
C     The inner bounding surface is a prolate spheroid with semi-axis
C     lengths
C
C        BMIN, BMIN, AMIN
C
C
C$ Examples
C
C     See usage in ZZRYXPDT.
C
C$ Restrictions
C
C     1)  The bounding spheroids generated by this routine may not be
C         suitable for reference spheroids with high eccentricity.
C
C     2)  Callers of this routine should not rely explicitly on the
C         formulas shown in the Detailed Output section of this
C         routine's header comments. The formulas could be upgraded
C         in a future version of this routine.
C
C$ Literature_References
C
C     None.
C
C$ Author_and_Institution
C
C     N.J. Bachman    (JPL)
C
C$ Version
C
C-    SPICELIB Version 1.0.0, 13-JAN-2017 (NJB)
C
C        Added error check for oblate case where lower bounding
C        ellipsoid has negative polar radius.
C
C        Updated $Particulars; added mention of use for volume elements
C        in planetodetic coordinate systems. Remarks about application
C        to prolate spheroids were added. Changed contents of $Files
C        section to "None." Made miscellaneous small changes to
C        comments.
C
C     Original version 27-NOV-2012 (NJB)
C
C-&

C$ Index_Entries
C
C     create bounding ellipsoids for a dem
C     bounding ellipsoids for a digital elevation model
C     create bounding ellipsoids for a dsk type 4 segment
C
C-&

C
C
C     Use discovery check-in.
C
      IF ( B .LE. 0.D0 ) THEN

         CALL CHKIN  ( 'ZZELLBDS'                                )
         CALL SETMSG ( 'This routine requires B > 0, but B = #.' )
         CALL ERRDP  ( '#',  B                                   )
         CALL SIGERR ( 'SPICE(NONPOSITIVERADIUS)'                )
         CALL CHKOUT ( 'ZZELLBDS'                                )
         RETURN

      END IF

      IF ( B .GT. A ) THEN

         CALL CHKIN  ( 'ZZELLBDS'                                 )
         CALL SETMSG ( 'This routine requires A >= B, but A = #; '
     .   //            'B = #.'                                   )
         CALL ERRDP  ( '#',  A                                    )
         CALL ERRDP  ( '#',  B                                    )
         CALL SIGERR ( 'SPICE(RADIIOUTOFORDER)'                   )
         CALL CHKOUT ( 'ZZELLBDS'                                 )
         RETURN

      END IF

      IF ( ( B + HMIN ) .LE. 0.D0 ) THEN

         CALL CHKIN  ( 'ZZELLBDS'                                 )
         CALL SETMSG ( 'This routine requires B + HMIN > 0, but '
     .   //            'B = #; HMIN = #, B+HMIN = #.'             )
         CALL ERRDP  ( '#',  B                                    )
         CALL ERRDP  ( '#',  HMIN                                 )
         CALL ERRDP  ( '#',  B + HMIN                             )
         CALL SIGERR ( 'SPICE(LOWERBOUNDTOOLOW)'                  )
         CALL CHKOUT ( 'ZZELLBDS'                                 )
         RETURN

      END IF

      IF ( HMIN .LT. 0.D0 ) THEN

         IF (  ( B + (A/B)*HMIN ) .LE. 0.D0 ) THEN

            CALL CHKIN  ( 'ZZELLBDS'                                 )
            CALL SETMSG ( 'For oblate spheroids and HMIN < 0, '
     .      //            'This routine requires B + (A/B)HMIN > 0, '
     .      //            'but A = #, B = #; HMIN = #, '
     .      //            'B+(A/B)HMIN = #.'                         )
            CALL ERRDP  ( '#',  A                                    )
            CALL ERRDP  ( '#',  B                                    )
            CALL ERRDP  ( '#',  HMIN                                 )
            CALL ERRDP  ( '#',  B + (A/B)*HMIN                       )
            CALL SIGERR ( 'SPICE(LOWERBOUNDTOOLOW)'                  )
            CALL CHKOUT ( 'ZZELLBDS'                                 )
            RETURN

         END IF

      END IF

      IF ( HMIN .GT. HMAX ) THEN

         CALL CHKIN  ( 'ZZELLBDS'                                 )
         CALL SETMSG ( 'This routine requires HMAX >= HMIN, but '
     .   //            'HMIN = #; HMAX = #.'                      )
         CALL ERRDP  ( '#',  HMIN                                 )
         CALL ERRDP  ( '#',  HMAX                                 )
         CALL SIGERR ( 'SPICE(BOUNDSOUTOFORDER)'                  )
         CALL CHKOUT ( 'ZZELLBDS'                                 )
         RETURN

      END IF

C
C     In the following comments, N, E, E', and LAMBDA are
C     defined as in the Particulars section above.
C
C
C     Generate radii of the outer bounding ellipsoid.
C
      IF ( HMAX .GE. 0.D0 ) THEN
C
C        Pick radii of E' so that E' matches
C
C           E + HMAX * N / ||N||
C
C        that is, E' has height HMAX above E, at x=A.
C
C        For smaller x, the height of E' above E will
C        will be greater than or equal to HMAX.
C
C        Set LAMBDA = A * HMAX.
C
C        Then the radii of E' are
C
C                             |
C           x + LAMBDA*||N||  |
C                             |x=A,y=0
C
C        and
C                             |
C           y + LAMBDA*||N||  |
C                             |x=0,y=B
C
C        so the radii of E', AMAX and BMAX, are:
C
C           AMAX =  A + LAMBDA*A/A**2  =  A + HMAX
C           BMAX =  B + LAMBDA*B/B**2  =  B + HMAX*(A/B)
C
C
         AMAX = A  +  HMAX
         BMAX = B  +  HMAX * (A/B)

      ELSE
C
C        HMAX < 0.
C
C        In this case the outer bounding ellipse should match E+HMAX
C        at x = 0. The ellipse will be closer to E for x > 0.
C
C        Set LAMBDA = B * HMAX. Then
C
C           AMAX =  A + LAMBDA*A/A**2  =  A + HMAX * (B/A)
C           BMAX =  B + LAMBDA*B/B**2  =  B + HMAX

         AMAX = A  +  HMAX * (B/A)
         BMAX = B  +  HMAX

      END IF

C
C     Find radii of the inner bounding ellipsoid.
C
      IF ( HMIN .LE. 0 ) THEN
C
C        This case is similar to that of the outer bounding
C        ellipsoid for HMAX >= 0. We can create an ellipse
C        that has height HMIN at x = A and that is further
C        from E for x < A.
C
C        Set LAMBDA = A * HMAX. Then
C
C           AMAX =  A + LAMBDA*A/A**2  =  A + HMAX
C           BMAX =  B + LAMBDA*B/B**2  =  B + HMAX*(A/B)

         AMIN = A  +  HMIN
         BMIN = B  +  HMIN * (A/B)

      ELSE
C
C        HMIN > 0.
C
C        In this case the inner bounding ellipse should match E+HMIN
C        at x = 0. The ellipse will be closer to E for x > 0.
C
C        Set LAMBDA = B * HMAX. Then
C
C           AMIN =  A + LAMBDA*A/A**2  =  A + HMIN * (B/A)
C           BMIN =  B + LAMBDA*B/B**2  =  B + HMIN
C
         AMIN = A  +  HMIN * (B/A)
         BMIN = B  +  HMIN

      END IF

      END


REAL (C_DOUBLE) FUNCTION ddot_noblas(n, x, y) BIND(C)
    USE ISO_C_BINDING
    IMPLICIT NONE

    INTEGER (C_LONG), VALUE :: n
    REAL (C_DOUBLE), DIMENSION(n) :: x, y

    ddot_noblas = DOT_PRODUCT(x, y)

END FUNCTION

REAL (C_DOUBLE) FUNCTION ddot_blas(n, x, y) BIND(C)
    USE ISO_C_BINDING
    IMPLICIT NONE

    INTEGER (C_LONG), VALUE :: n
    REAL (C_DOUBLE), DIMENSION(n) :: x, y
    DOUBLE PRECISION ddot

    ddot_blas = ddot(n, x, 1, y, 1)

END FUNCTION


        !COMPILER-GENERATED INTERFACE MODULE: Fri May 22 17:00:09 2020
        ! This source file is for reference only and may not completely
        ! represent the generated interface used by the compiler.
        MODULE SMOOTHING_WEIGHTS_VEC__genmod
          INTERFACE
            SUBROUTINE SMOOTHING_WEIGHTS_VEC(X0,Z0,SIGMA_H2_INV,        &
     &SIGMA_V2_INV,EXP_VAL,XX_ELEM,ZZ_ELEM)
              REAL(KIND=4), INTENT(IN) :: X0
              REAL(KIND=4), INTENT(IN) :: Z0
              REAL(KIND=4), INTENT(IN) :: SIGMA_H2_INV
              REAL(KIND=4), INTENT(IN) :: SIGMA_V2_INV
              REAL(KIND=4), INTENT(OUT) :: EXP_VAL(5,5)
              REAL(KIND=4), INTENT(IN) :: XX_ELEM(5,5)
              REAL(KIND=4), INTENT(IN) :: ZZ_ELEM(5,5)
            END SUBROUTINE SMOOTHING_WEIGHTS_VEC
          END INTERFACE
        END MODULE SMOOTHING_WEIGHTS_VEC__genmod


!=======================================================================
! Created by
!     Keurfon Luu <keurfon.luu@mines-paristech.fr>
!     MINES ParisTech - Centre de Géosciences
!     PSL - Research University
!=======================================================================

program sampling

  use forlab, only: IPRE, RPRE, CLEN, split_argument, ones, num2str, rng
  use benchmark_functions, only: BenchmarkFunction
  use stochoptim, only: MonteCarlo
#ifdef do_mpi
  use mpi
#endif

  implicit none

  ! Inputs
  integer(kind = IPRE) :: n_dim, max_iter, seed, verbose
  real(kind = RPRE) :: stepsize, perc
  character(len = :), allocatable :: sampler, func

  ! Local variables
  character(len = 18) :: attributes
  type(MonteCarlo) :: mc
  type(BenchmarkFunction) :: bf

  ! MPI variables
  integer(kind = IPRE) :: master = 0, rank, ierr, iprv

  ! Initialize MPI
#ifdef do_mpi
  call mpi_init_thread(mpi_thread_funneled, iprv, ierr)
  call mpi_comm_rank(mpi_comm_world, rank, ierr)
#else
  rank = 0
#endif

  ! Read command line arguments
  call command_arguments()

  ! Initialize random number generator
  if ( seed .ge. 0 ) then
    call rng(seed)
  else
    call rng()
  end if

  bf = BenchmarkFunction(func, n_dim)
  mc = MonteCarlo(bf % func, bf % lower, bf % upper, max_iter = max_iter, &
                  perc = perc, constrain = .true., verbose = verbose)

  call mc % sample(&
    sampler = sampler, &
    stepsize = stepsize * ones(n_dim))

  ! Display results
  if ( rank .eq. master ) then
    print *
    attributes = "function:"
    print *, adjustr(attributes) // " '" // trim(func) // "'"
    call mc % print()
    print *
  end if

  ! Save results
  if ( rank .eq. master ) then
    call system(&
      "rm -rf output;" // &
      "mkdir -p output")
    call mc % save_models("output/")
  end if

#ifdef do_mpi
  call mpi_finalize(ierr)
#endif

contains

  subroutine command_arguments()
    integer(kind = IPRE) :: i, n_argin
    character(len = CLEN) :: argin
    character(len = :), allocatable :: argname, argval

    sampler = "mcmc"
    func = "sphere"
    n_dim = 2
    max_iter = 1000
    stepsize = 0.1
    perc = 1.
    seed = -1
    verbose = 0

    n_argin = command_argument_count()
    do i = 1, n_argin
      call get_command_argument(i, argin)
      call split_argument(argin, argname, argval)

      select case(argname)
      ! Number of dimensions
      case("n_dim")
        read(argval, *) n_dim

      ! Sampler
      case("sampler")
        sampler = argval

      ! Function
      case("func")
        func = argval

      ! Maximum number of iterations
      case("max_iter")
        read(argval, *) max_iter

      ! Stepsize
      case("stepsize")
        read(argval, *) stepsize

      ! Number of dimensions to perturbe as a percentage
      case("perc")
        read(argval, *) perc

      ! Random number seed
      case("seed")
        read(argval, *) seed

      ! Verbosity
      case("verbose")
        read(argval, *) verbose

      case default
        print *, "Error: unknown argument '" // trim(argname) // "'"
        stop
      end select
    end do
    return
  end subroutine command_arguments

end program sampling


! ##################################################################################################################################
! Begin MIT license text.
! _______________________________________________________________________________________________________

! Copyright 2019 Dr William R Case, Jr (dbcase29@gmail.com)

! Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
! associated documentation files (the "Software"), to deal in the Software without restriction, including
! without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
! copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to
! the following conditions:

! The above copyright notice and this permission notice shall be included in all copies or substantial
! portions of the Software and documentation.

! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
! OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
! FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
! AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
! LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
! OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
! THE SOFTWARE.
! _______________________________________________________________________________________________________

! End MIT license text.

      SUBROUTINE SOLVE_DLR

! Solves KLL*DLR = -KLR for matrix DLR. However, we will use rows of KRL instead of cols of KLR in the solution.

! For a description of Craig-Bamptom analyses, see Appendix D to the MYSTRAN User's Referance Manual

      USE PENTIUM_II_KIND, ONLY       :  BYTE, LONG, DOUBLE
      USE IOUNT1, ONLY                :  FILE_NAM_MAXLEN, WRT_ERR, WRT_LOG, ERR, F04, F06, SCR
      USE SCONTR, ONLY                :  BLNK_SUB_NAM, FACTORED_MATRIX, FATAL_ERR, KLL_SDIA, NDOFR, NDOFL, NTERM_DLR, NTERM_KLL,   &
                                         NTERM_KRL
      USE PARAMS, ONLY                :  EPSIL, PRTDLR, SOLLIB, SPARSE_FLAVOR, SPARSTOR
      USE TIMDAT, ONLY                :  HOUR, MINUTE, SEC, SFRAC, TSEC
      USE SUBR_BEGEND_LEVELS, ONLY    :  SOLVE_DLR_BEGEND
      USE CONSTANTS_1, ONLY           :  ZERO, ONE
      USE SPARSE_MATRICES, ONLY       :  I2_DLR, I_DLR, J_DLR, DLR, I_DLRt, I2_DLRt, J_DLRt, DLRt, I_KRL, J_KRL, KRL,              &
                                         I_KLL, I2_KLL, J_KLL, KLL

      USE LAPACK_LIN_EQN_DPB

      USE SOLVE_DLR_USE_IFs

      IMPLICIT NONE

      CHARACTER, PARAMETER            :: CR13 = CHAR(13)   ! This causes a carriage return simulating the "+" action in a FORMAT
      CHARACTER(LEN=LEN(BLNK_SUB_NAM)):: SUBR_NAME = 'SOLVE_DLR'
      CHARACTER(  1*BYTE)             :: CLOSE_IT          ! Input to subr READ_MATRIX_i. 'Y'/'N' whether to close a file or not
      CHARACTER(  8*BYTE)             :: CLOSE_STAT        ! What to do with file when it is closed
      CHARACTER(  1*BYTE)             :: EQUED             ! 'Y' if KLL stiff matrix was equilibrated in subr EQUILIBRATE
      CHARACTER( 24*BYTE)             :: MESSAG            ! File description. Input to subr UNFORMATTED_OPEN
      CHARACTER( 22*BYTE)             :: MODNAM1           ! Name to write to screen to describe module being run
      CHARACTER(  1*BYTE)             :: READ_NTERM        ! 'Y' or 'N' Input to subr READ_MATRIX_1
      CHARACTER(  1*BYTE)             :: NULL_COL          ! 'Y' if a col of KLR(transpose) is null
      CHARACTER(  1*BYTE)             :: OPND              ! Input to subr READ_MATRIX_i. 'Y'/'N' whether to open  a file or not
      CHARACTER(FILE_NAM_MAXLEN*BYTE) :: SCRFIL            ! File name

      INTEGER(LONG)                   :: DEB_PRT(2)        ! Debug numbers to say whether to write ABAND and/or its decomp to output
!                                                            file in called subr SYM_MAT_DECOMP_LAPACK
      INTEGER(LONG)                   :: I,J               ! DO loop indices or counters
      INTEGER(LONG)                   :: INFO        = 0   ! Info on success of factorization or solve
      INTEGER(LONG)                   :: IOCHK             ! IOSTAT error number when opening a file
      INTEGER(LONG)                   :: OUNT(2)           ! File units to write messages to. Input to subr UNFORMATTED_OPEN
      INTEGER(LONG), PARAMETER        :: SUBR_BEGEND = SOLVE_DLR_BEGEND

      REAL(DOUBLE)                    :: EPS1              ! A small number to compare real zero

      REAL(DOUBLE)                    :: EQUIL_SCALE_FACS(NDOFL)
                                                           ! LAPACK_S values returned from subr SYM_MAT_DECOMP_LAPACK

      REAL(DOUBLE)                    :: DLR_COL(NDOFL)    ! A column of DLR solved for herein
      REAL(DOUBLE)                    :: INOUT_COL(NDOFL)  ! Temp variable for subr FBS
      REAL(DOUBLE)                    :: K_INORM           ! Inf norm of KLL matrix (det in  subr COND_NUM)
      REAL(DOUBLE)                    :: RCOND             ! Recrip of cond no. of the KLL. Det in  subr COND_NUM

      INTRINSIC                       :: DABS

! **********************************************************************************************************************************
      IF (WRT_LOG >= SUBR_BEGEND) THEN
         CALL OURTIM
         WRITE(F04,9001) SUBR_NAME,TSEC
 9001    FORMAT(1X,A,' BEGN ',F10.3)
      ENDIF

! **********************************************************************************************************************************
      EPS1 = EPSIL(1)

! Decomp KLL

      DEB_PRT(1) = 64
      DEB_PRT(2) = 65

      IF (SOLLIB == 'BANDED') THEN

         INFO  = 0
         EQUED = 'N'
         CALL SYM_MAT_DECOMP_LAPACK ( SUBR_NAME, 'KLL', 'L ', NDOFL, NTERM_KLL, I_KLL, J_KLL, KLL, 'Y', 'N', 'N', 'N', DEB_PRT, &
                                      EQUED, KLL_SDIA, K_INORM, RCOND, EQUIL_SCALE_FACS, INFO )
         IF (EQUED == 'Y') THEN                         ! If EQUED == 'Y' then error. We don't want KLL equilibrated
            WRITE(ERR,6001) SUBR_NAME, EQUED
            WRITE(F06,6001) SUBR_NAME, EQUED
            FATAL_ERR = FATAL_ERR + 1
            CALL OUTA_HERE ( 'Y' )
         ENDIF

      ELSE IF (SOLLIB == 'SPARSE  ') THEN

         IF (SPARSE_FLAVOR(1:7) == 'SUPERLU') THEN

            INFO = 0
            CALL SYM_MAT_DECOMP_SUPRLU ( SUBR_NAME, 'KLL', NDOFL, NTERM_KLL, I_KLL, J_KLL, KLL, INFO )

         ELSE

            FATAL_ERR = FATAL_ERR + 1
            WRITE(ERR,9991) SUBR_NAME, 'SPARSE_FLAVOR'
            WRITE(F06,9991) SUBR_NAME, 'SPARSE_FLAVOR'
            CALL OUTA_HERE ( 'Y' )

         ENDIF

      ELSE

         FATAL_ERR = FATAL_ERR + 1
         WRITE(ERR,9991) SUBR_NAME, 'SOLLIB'
         WRITE(F06,9991) SUBR_NAME, 'SOLLIB'
         CALL OUTA_HERE ( 'Y' )

      ENDIF

      DO I=1,NDOFL                                         ! Make sure that scale factors are one. We don't want any equil scaling
         EQUIL_SCALE_FACS(I) = ONE                         ! of KLL in this subr. FBS below has EQUIL_SCALE_FACS as an input but
      ENDDO                                                ! they shouldn't be used as EQUED = 'N' is also input there (1st arg)

!***********************************************************************************************************************************
! Solve for DLR

! Open a scratch file that will be used to write DLR nonzero terms to as we solve for columns of DLR. After all col's
! of DLR have been solved for, and we have a count on NTERM_DLR, we will allocate memory to the DLR arrays and read
! the scratch file values into those arrays. Then, in the calling subroutine, we will write NTERM_DLR, followed by
! DLR  row/col/value to a permanent file

      OUNT(1) = ERR
      OUNT(2) = F06
      SCRFIL(1:)  = ' '
      SCRFIL(1:9) = 'SCRATCH-991'
      OPEN (SCR(1),STATUS='SCRATCH',FORM='UNFORMATTED',ACTION='READWRITE',IOSTAT=IOCHK)
      IF (IOCHK /= 0) THEN
         CALL OPNERR ( IOCHK, SCRFIL, OUNT, 'Y' )
         CALL FILE_CLOSE ( SCR(1), SCRFIL, 'DELETE', 'Y' )
         CALL OUTA_HERE ( 'Y' )                            ! Can't open scratch file, so quit
      ENDIF

! Loop on columns of KLR using rows of KRL

!xx   WRITE(SC1, * )                                       ! Advance 1 line for screen messages

      NTERM_DLR = 0
      DO J = 1,NDOFR

         CALL OURTIM
         MODNAM1 = '    Solve for DLR col '
         WRITE(SC1,12345,ADVANCE='NO') MODNAM1, J, NDOFR, CR13

! To solve for the j-th col of DLR, use the j-th row of KRL (j-th col of KLR) as a rhs vector. Get the j-th row of KRL and put
! the negative of it into array INOUT_COL:

         NULL_COL = 'Y'
         DO I=1,NDOFL
            INOUT_COL(I) = ZERO
            DLR_COL(I)   = ZERO
         ENDDO
         CALL GET_SPARSE_CRS_ROW ( 'KRL', J,  NTERM_KRL, NDOFR, NDOFL, I_KRL, J_KRL, KRL, -ONE, INOUT_COL, NULL_COL )

! Calculate DLR_COL via forward/backward substitution.

         IF (NULL_COL == 'N') THEN                         ! FBS will solve for DLR_COL & load it into DLR array
                                                           ! DPBTRS will return DLR_COL = -KLL(-1)*RHS_col
!                                                            Note 1st arg = 'N' assures that EQUIL_SCAL_FACS will not be used
            IF      (SOLLIB == 'BANDED') THEN

               CALL FBS_LAPACK ( 'N', NDOFL, KLL_SDIA, EQUIL_SCALE_FACS, INOUT_COL )

            ELSE IF (SOLLIB == 'SPARSE  ') THEN

               IF (SPARSE_FLAVOR(1:7) == 'SUPERLU') THEN

                  INFO = 0
                  CALL FBS_SUPRLU ( SUBR_NAME, 'KLL', NDOFL, NTERM_KLL, I_KLL, J_KLL, KLL, J, INOUT_COL, INFO )

               ELSE

                  FATAL_ERR = FATAL_ERR + 1
                  WRITE(ERR,9991) SUBR_NAME, 'SPARSE_FLAVOR'
                  WRITE(F06,9991) SUBR_NAME, 'SPARSE_FLAVOR'
                  CALL OUTA_HERE ( 'Y' )

               ENDIF

            ELSE

               FATAL_ERR = FATAL_ERR + 1
               WRITE(ERR,9991) SUBR_NAME, 'SOLLIB'
               WRITE(F06,9991) SUBR_NAME, 'SOLLIB'
               CALL OUTA_HERE ( 'Y' )

            ENDIF

            DO I=1,NDOFL
               DLR_COL(I) = INOUT_COL(I)
            ENDDO
            DO I=1,NDOFL                                   ! Count NTERM_DLR and write nonzero DLR to scratch file
               IF (DABS(DLR_COL(I)) > EPS1) THEN
                  NTERM_DLR = NTERM_DLR + 1
                  WRITE(SCR(1)) I, J, DLR_COL(I)
               ENDIF
            ENDDO
         ENDIF

      ENDDO

! The DLR data in SCRATCH-991 is written one col at a time for DLR. Therefore it is rows of DLRt

      REWIND (SCR(1))
      MESSAG = 'SCRATCH: DLR ROW/COL/VAL'
      READ_NTERM = 'N'
      OPND       = 'Y'
      CLOSE_IT   = 'N'
      CLOSE_STAT = 'KEEP    '

      CALL ALLOCATE_SPARSE_MAT ( 'DLRt', NDOFR, NTERM_DLR, SUBR_NAME )
      CALL ALLOCATE_L6_2 ( 'DLRt', SUBR_NAME )

      CALL ALLOCATE_SPARSE_MAT ( 'DLR', NDOFL, NTERM_DLR, SUBR_NAME )
      CALL ALLOCATE_L6_2 ( 'DLR', SUBR_NAME )
                                                           ! J_DLRt is same as I2_DLR and I2_DLRt is same as J_DLR
      CALL READ_MATRIX_2 ( SCRFIL, SCR(1), OPND, CLOSE_IT, CLOSE_STAT, MESSAG, 'DLRt', NDOFR, NTERM_DLR, READ_NTERM,               &
                           J_DLRt, I2_DLRt, DLRt)

! Now get DLR from DLRt

      CALL GET_I_MAT_FROM_I2_MAT ( 'DLRt', NDOFR, NTERM_DLR, I2_DLRt, I_DLRt )

      CALL MATTRNSP_SS ( NDOFR, NDOFL, NTERM_DLR, 'DLRt', I_DLRt, J_DLRt, DLRt, 'DLR', I_DLR, J_DLR, DLR )

      CALL FILE_CLOSE ( SCR(1), SCRFIL, 'DELETE', 'Y' )

! Print out constraint matrix DLR, if requested

      IF ( PRTDLR == 1) THEN
         IF (NTERM_DLR > 0) THEN
            CALL WRITE_SPARSE_CRS ( 'CB BOUNDARY MODE MATRIX DLR', 'L ', 'R ', NTERM_DLR, NDOFL, I_DLR, J_DLR, DLR )
         ENDIF
      ENDIF

! **********************************************************************************************************************************
      IF (WRT_LOG >= SUBR_BEGEND) THEN
         CALL OURTIM
         WRITE(F04,9002) SUBR_NAME,TSEC
 9002    FORMAT(1X,A,' END  ',F10.3)
      ENDIF

      RETURN

! **********************************************************************************************************************************
  932 FORMAT(' *ERROR   932: PROGRAMMING ERROR IN SUBROUTINE ',A                                                                   &
                    ,/,14X,' PARAMETER SPARSTOR MUST BE EITHER "SYM" OR "NONSYM" BUT VALUE IS ',A)

 2092 FORMAT(4X,A44,20X,I2,':',I2,':',I2,'.',I3)

 6001 FORMAT(' *ERROR  6001: PROGRAMMING ERROR IN SUBROUTINE ',A                                                                   &
                    ,/,14X,' MATRIX KLL WAS EQUILIBRATED: EQUED = ',A,'. CODE NOT WRITTEN TO ALLOW THIS AS YET')

 9991 FORMAT(' *ERROR  9991: PROGRAMMING ERROR IN SUBROUTINE ',A                                                                   &
                    ,/,14X,A, ' = ',A,' NOT PROGRAMMED ',A)

12345 FORMAT(3X,A,I8,' of ',I8,A)

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

      END SUBROUTINE SOLVE_DLR


module gyro_averages

   use common_types, only: coupled_alpha_type

   public :: aj0x, aj0v, aj1x, aj1v
   public :: init_bessel, finish_bessel
   public :: gyro_average
   public :: gyro_average_j1
   public :: j0_B_maxwell_ffs, j0_ffs
   public :: band_lu_solve_ffs, band_lu_factorisation_ffs

   private

   interface gyro_average
      module procedure gyro_average_kxky_local
      module procedure gyro_average_kxkyz_local
      module procedure gyro_average_vmu_local
      module procedure gyro_average_vmus_nonlocal
      module procedure gyro_average_ffs_kxky_local
      module procedure gyro_average_ffs_kxkyz_local
      module procedure gyro_average_ffs
   end interface gyro_average

   interface gyro_average_j1
      module procedure gyro_average_j1_kxky_local
      module procedure gyro_average_j1_kxkyz_local
      module procedure gyro_average_j1_vmu_local
   end interface

   real, dimension(:, :, :, :), allocatable :: aj0x, aj1x
   ! (naky, nakx, nalpha, -nzgrid:nzgrid, -vmu-layout-)

   real, dimension(:, :), allocatable :: aj0v, aj1v
   ! (nmu, -kxkyz-layout-)

   type(coupled_alpha_type), dimension(:, :, :, :), allocatable :: j0_ffs, j0_B_maxwell_ffs

   logical :: bessinit = .false.

   logical :: debug = .false.

contains

   subroutine init_bessel

      use dist_fn_arrays, only: kperp2
      use physics_flags, only: full_flux_surface
      use species, only: spec, nspec
      use stella_geometry, only: bmag
      use zgrid, only: nzgrid
      use vpamu_grids, only: vperp2, nmu
      use kt_grids, only: naky, nakx
      use stella_layouts, only: kxkyz_lo, vmu_lo
      use stella_layouts, only: iky_idx, ikx_idx, iz_idx, is_idx, imu_idx
      use spfunc, only: j0, j1

      implicit none

      integer :: iz, iky, ikx, imu, is, ia
      integer :: ikxkyz, ivmu
      real :: arg

      if (bessinit) return
      bessinit = .true.

      if (debug) write (*, *) 'gyro_averages::init_bessel::allocate_aj0v_aj1v'
      if (.not. allocated(aj0v)) then
         allocate (aj0v(nmu, kxkyz_lo%llim_proc:kxkyz_lo%ulim_alloc))
         aj0v = 0.
      end if
      if (.not. allocated(aj1v)) then
         allocate (aj1v(nmu, kxkyz_lo%llim_proc:kxkyz_lo%ulim_alloc))
         aj1v = 0.
      end if

      if (debug) write (*, *) 'gyro_averages::init_bessel::calculate_aj0v_aj1v'
      ia = 1
      do ikxkyz = kxkyz_lo%llim_proc, kxkyz_lo%ulim_proc
         iky = iky_idx(kxkyz_lo, ikxkyz)
         ikx = ikx_idx(kxkyz_lo, ikxkyz)
         iz = iz_idx(kxkyz_lo, ikxkyz)
         is = is_idx(kxkyz_lo, ikxkyz)
         do imu = 1, nmu
            arg = spec(is)%bess_fac * spec(is)%smz_psi0 * sqrt(vperp2(ia, iz, imu) * kperp2(iky, ikx, ia, iz)) / bmag(ia, iz)
            aj0v(imu, ikxkyz) = j0(arg)
            ! note that j1 returns and aj1 stores J_1(x)/x (NOT J_1(x)),
            aj1v(imu, ikxkyz) = j1(arg)
         end do
      end do

      if (debug) write (*, *) 'gyro_averages::init_bessel::full_flux_surface'
      if (full_flux_surface) then
         call init_bessel_ffs
      else
         if (.not. allocated(aj0x)) then
            allocate (aj0x(naky, nakx, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
            aj0x = 0.
         end if

         if (.not. allocated(aj1x)) then
            allocate (aj1x(naky, nakx, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
            aj1x = 0.
         end if

         ia = 1
         do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
            is = is_idx(vmu_lo, ivmu)
            imu = imu_idx(vmu_lo, ivmu)
            do iz = -nzgrid, nzgrid
               do ikx = 1, nakx
                  do iky = 1, naky
                     arg = spec(is)%bess_fac * spec(is)%smz_psi0 * sqrt(vperp2(ia, iz, imu) * kperp2(iky, ikx, ia, iz)) / bmag(ia, iz)
                     aj0x(iky, ikx, iz, ivmu) = j0(arg)
                     ! note that j1 returns and aj1 stores J_1(x)/x (NOT J_1(x)),
                     aj1x(iky, ikx, iz, ivmu) = j1(arg)
                  end do
               end do
            end do
         end do
      end if
      if (debug) write (*, *) 'gyro_averages::init_bessel::test_gyro_average'
!    call test_gyro_average

   contains

!     ! set up field that varies as x^2 = rho^2 * cos(angle)^2,
!     ! with rho the distance from the origin, and 'angle' is the angle made with the horizontal
!     ! if considering a particle at x=0, then rho is thee gyro-radius and angle is the gyro-angle
!     ! the gyro-average should theen be 1/(2pi) * int_0^2pi dangle rho^2 * cos(angle)^2 = rho^2/2
!     subroutine test_gyro_average

!       use constants, only: pi
!       use kt_grids, only: ny, nx, x, x0, y, y0
!       use kt_grids, only: nakx, ikx_max, naky, naky_all
!       use kt_grids, only: swap_kxky, swap_kxky_back
!       use stella_transforms, only: transform_x2kx, transform_y2ky
!       use stella_transforms, only: transform_kx2x, transform_ky2y
!       use stella_layouts, only: vmu_lo, iv_idx, imu_idx, is_idx
!       use vpamu_grids, only: nmu
!       use species, only: nspec, spec
!       use stella_geometry, only: alpha, bmag, x_displacement_fac
!       use spfunc, only: bessi0

!       implicit none

!       real, dimension (:,:), allocatable :: fld_yx
!       complex, dimension (:,:), allocatable :: fld_ykx
!       complex, dimension (:,:), allocatable :: fld_kykx_swapped
!       complex, dimension (:,:), allocatable :: fld_kykx, gyro_fld
!       real, dimension (:,:,:,:), allocatable :: gyro_fld_yx
!       real :: gyroradius
!       integer :: iy, ix, ivmu, iv, imu, is

!       integer, parameter :: iz = 0

!       allocate (fld_yx(ny,nx))
!       allocate (fld_ykx(ny,ikx_max))
!       allocate (fld_kykx_swapped(naky_all,ikx_max))
!       allocate (fld_kykx(naky,nakx))
!       allocate (gyro_fld(naky,nakx))
!       allocate (gyro_fld_yx(ny,nx,nmu,nspec))

!       if (debug) write (*,*) 'gyro_averages::init_bessel::test_gyro_averages::fld_yx'
!       ! set up field that varies as x^2 = rho^2 * cos(angle)^2 and is constant in y
! !      fld_yx = spread(0.1*(x-pi*x0),1,ny)**2
!       fld_yx = spread(cos(50.0*(x/x0-pi)),1,ny)
! !      fld_yx = spread(exp(-0.1*(x-pi*x0)**2),1,ny)
!       if (debug) write (*,*) 'gyro_averages::init_bessel::test_gyro_averages::transform_x2kx'
!       ! transform from (y,x) to (y,kx), with kx going from 0 to kxmax
!       call transform_x2kx (fld_yx, fld_ykx)
!       if (debug) write (*,*) 'gyro_averages::init_bessel::test_gyro_averages::transform_y2ky'
!       ! transform from (y,kx) to (ky,kx), with ky going from (0,...,kymax,-kymax,...,-dky)
!       ! and kx going from 0 to kxmax
!       call transform_y2ky (fld_ykx, fld_kykx_swapped)
!       if (debug) write (*,*) 'gyro_averages::init_bessel::test_gyro_averages::swap_kxky_back'
!       ! use reality condition to re-arrange array so that ky goes from 0 to kymax
!       ! and kx goes from (0,...,kxmax,-kxmax,...,-dkx)
!       call swap_kxky_back (fld_kykx_swapped, fld_kykx)

!       ! gyro-average the field at z=0 for different values of mu
!       do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
!          ! get the vpa index
!          iv = iv_idx(vmu_lo,ivmu)
!          ! as J0 independent of vpa, pick only one vpa to test
!          if (iv /= 1) cycle
!          ! get the species index
!          is = is_idx(vmu_lo,ivmu)
!          ! get the mu index
!          imu = imu_idx(vmu_lo,ivmu)
! !         if (debug) write (*,*) 'gyro_averages::init_bessel::test_gyro_averages::gyro_average'
!          if (full_flux_surface) then
!             call gyro_average (fld_kykx, gyro_fld, j0_ffs(:,:,iz,ivmu))
!          else
!             call gyro_average (fld_kykx, iz, ivmu, gyro_fld)
!          end if
!          ! use reality to re-arrange array entries so that ky goes from (0,...,kymax,-kymax,...,-dky)
!          ! and kx goes from 0 to kxmax
!          call swap_kxky (gyro_fld, fld_kykx_swapped)
!          ! transform from (ky,kx) to (y,kx)
!          call transform_ky2y (fld_kykx_swapped, fld_ykx)
!          ! transform from (y,kx) to (y,x)
!          call transform_kx2x (fld_ykx, gyro_fld_yx(:,:,imu,is))
!       end do
!       if (debug) write (*,*) 'gyro_averages::init_bessel::test_gyro_averages::write_to_screen'
!       ! NB: this is only set up to work on a single processor at the moment
!       ! NB: to extend, must move information about gyro_fld onto proc0
!       do is = 1, nspec
!          do imu = 1, nmu
!             do ix = 1, nx
!                do iy = 1, ny
!                   ! gyro-radius/reference gyro-radius is v_perp/Omega/rho_ref = (v_perp/vths)*(rho_s/rho_ref)
!                   ! = vperp * sqrt(T_s/T_ref*m_s/m_ref)(B_ref/Z*B) = vperp / (spec%zstm*bmag)
!                   gyroradius = sqrt(vperp2(iy,iz,imu))/(spec(is)%zstm*bmag(iy,iz))
! !                  gyroradius = sqrt(vperp2(1,iz,imu))/(spec(is)%zstm*bmag(1,iz))
!                   write (42,*) 'y: ', y(iy), 'x: ', x(ix)-x0*pi, 'gyro_fld: ', gyro_fld_yx(iy,ix,imu,is), 'gyroradius: ', gyroradius, 'spec: ', is, &
!                        'alpha: ', alpha(iy), 'x_displacement_fac: ', x_displacement_fac(iy,iz)
!                end do
!                write (42,*)
!             end do
!             write (42,*)
!          end do
!          write (42,*)
!       end do
!       do iy = 1, 1000
!          gyroradius = (iy-1)*15.0/999.
!          !         write (43,*) 'gyroradius: ', gyroradius, 'bes: ', bessi0(0.1*0.5*(gyroradius/x_displacement_fac(1,iz))**2)*exp(-0.1*0.5*(gyroradius/x_displacement_fac(1,iz))**2)
!          !         write (43,*) 'gyroradius: ', gyroradius, 'analytical: ', 0.5*(0.1*gyroradius/x_displacement_fac(1,iz))**2
!          write (43,*) 'gyroradius: ', gyroradius, 'analytical: ', j0(50.*gyroradius/(x0*x_displacement_fac(1,iz)))
!       end do
!       ! TMP FOR TESTING
! !      stop

!       deallocate (fld_yx, fld_ykx)
!       deallocate (fld_kykx_swapped, fld_kykx)
!       deallocate (gyro_fld, gyro_fld_yx)

!     end subroutine test_gyro_average

   end subroutine init_bessel

   subroutine init_bessel_ffs

      use mp, only: sum_allreduce, proc0
      use spfunc, only: j0
      use stella_layouts, only: vmu_lo
      use stella_layouts, only: iv_idx, imu_idx, is_idx
      use stella_transforms, only: transform_alpha2kalpha
      use species, only: nspec, spec
      use stella_geometry, only: bmag
      use zgrid, only: nzgrid, nztot
      use vpamu_grids, only: nmu, nvpa
      use vpamu_grids, only: vperp2, maxwell_vpa, maxwell_mu
      use kt_grids, only: nalpha, naky, naky_all, ikx_max
      use kt_grids, only: swap_kxky_ordered
      use dist_fn_arrays, only: kperp2

      implicit none

      !    integer :: j0_ffs_unit, j0_B_maxwell_ffs_unit
      integer :: iky, ikx, ia, iz
      integer :: ivmu, iv, imu, is
      integer :: ia_max_j0_count, ia_max_j0_B_maxwell_count
      real :: arg, rtmp
      real :: ia_max_j0_reduction_factor, ia_max_j0_B_maxwell_reduction_factor
      real, dimension(:), allocatable :: wgts
      real, dimension(:), allocatable :: aj0_alpha, j0_B_maxwell
      real, dimension(:, :, :), allocatable :: kperp2_swap
      complex, dimension(:), allocatable :: aj0_kalpha, j0_B_maxwell_kalpha

      !       call open_output_file (j0_ffs_unit, '.j0_ffs')
      !       call open_output_file (j0_B_maxwell_ffs_unit, '.j0_over_B_ffs')

      ! wgts are species-dependent factors appearing in Gamma0 factor
      allocate (wgts(nspec))
      wgts = spec%dens * spec%z**2 / spec%temp

      if (debug) write (*, *) 'gyro_averages::init_bessel::full_flux_surface::allocate_arrays'
      !> aj0_alpha will contain J_0 as a function of k_alpha and alpha
      allocate (aj0_alpha(nalpha))
      allocate (aj0_kalpha(naky))
      !> j0_B_maxwell will contain J_0*B*exp(-v^2) as a function of k_alpha and alpha
      allocate (j0_B_maxwell(nalpha))
      allocate (j0_B_maxwell_kalpha(naky))
      allocate (kperp2_swap(naky_all, ikx_max, nalpha))
      if (.not. allocated(j0_ffs)) then
         allocate (j0_ffs(naky_all, ikx_max, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
      end if
      if (.not. allocated(j0_B_maxwell_ffs)) then
         allocate (j0_B_maxwell_ffs(naky_all, ikx_max, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
      end if

      ia_max_j0_count = 0; ia_max_j0_B_maxwell_count = 0
      do iz = -nzgrid, nzgrid
         if (proc0) write (*, *) 'calculating Fourier coefficients needed for gyro-averaging with alpha variation; zed index: ', iz
         !> for each value of alpha, take kperp^2 calculated on domain kx = [-kx_max, kx_max] and ky = [0, ky_max]
         !> and use symmetry to obtain kperp^2 on domain kx = [0, kx_max] and ky = [-ky_max, ky_max]
         !> this makes later convolutions involving sums over all ky more straightforward
         if (debug) write (*, *) 'gyro_averages::init_bessel::full_flux_surface::swap_kxky'
         do ia = 1, nalpha
            call swap_kxky_ordered(kperp2(:, :, ia, iz), kperp2_swap(:, :, ia))
         end do
         if (debug) write (*, *) 'gyro_averages::init_bessel::full_flux_surface::j0_loop'
         do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
            is = is_idx(vmu_lo, ivmu)
            iv = iv_idx(vmu_lo, ivmu)
            imu = imu_idx(vmu_lo, ivmu)
            do ikx = 1, ikx_max
               do iky = 1, naky_all
                  do ia = 1, nalpha
                     !> calculate the argument of the Bessel function, which depends on both alpha and k_alpha
                     arg = spec(is)%bess_fac * spec(is)%smz_psi0 * sqrt(vperp2(ia, iz, imu) * kperp2_swap(iky, ikx, ia)) / bmag(ia, iz)
                     ! compute the value of the Bessel function J0 corresponding to argument arg
                     aj0_alpha(ia) = j0(arg)
                     !> compute J_0*B*exp(-v^2), needed when integrating g over v-space in Maxwell's equations,
                     !> due to B in v-space Jacobian and Maxwellian factor hidden in normalisation of g
                     j0_B_maxwell(ia) = aj0_alpha(ia) * bmag(ia, iz) * maxwell_vpa(iv, is) * maxwell_mu(ia, iz, imu, is)
                  end do
                  !> fourier transform aj0_alpha and j0_B_maxwell.
                  !> note that fourier coefficients aj0_kalpha and j0_B_maxwell_kalpha have
                  !> been filtered to avoid aliasing
                  call transform_alpha2kalpha(aj0_alpha, aj0_kalpha)
                  call transform_alpha2kalpha(j0_B_maxwell, j0_B_maxwell_kalpha)
                  !> given the Fourier coefficients aj0_kalpha, calculate the minimum number of coefficients needed,
                  !> called j0_ffs%max_idx, to ensure that the relative error in the total spectral energy is below a specified tolerance
                  !if (debug) write (*,*) 'gyro_averages::init_bessel::full_flux_surface::find_max_required_kalpha_index'
                  !                ! TMP FOR TESTING
                  !                j0_ffs(iky,ikx,iz,ivmu)%max_idx = naky
                  call find_max_required_kalpha_index(aj0_kalpha, j0_ffs(iky, ikx, iz, ivmu)%max_idx, imu, iz, is)
                  !> given the Fourier coefficients j0_B_maxwell_kalpha, calculate the minimum number of coefficients needed,
                  !> called j0_B_maxwell_ffs%max_idx, to ensure that the relative error in the total spectral energy is below a specified tolerance
                  call find_max_required_kalpha_index(j0_B_maxwell_kalpha, j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%max_idx, imu, iz, is)
                  !> keep track of the total number of coefficients that must be retained across different phase space points
                  ia_max_j0_count = ia_max_j0_count + j0_ffs(iky, ikx, iz, ivmu)%max_idx
                  !> keep track of the total number of coefficients that must be retained across different phase space points
                  ia_max_j0_B_maxwell_count = ia_max_j0_B_maxwell_count + j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%max_idx
                  !> allocate array to hold the reduced number of Fourier coefficients
                  if (.not. associated(j0_ffs(iky, ikx, iz, ivmu)%fourier)) &
                     allocate (j0_ffs(iky, ikx, iz, ivmu)%fourier(j0_ffs(iky, ikx, iz, ivmu)%max_idx))
                  !> fill the array with the requisite coefficients
                  j0_ffs(iky, ikx, iz, ivmu)%fourier = aj0_kalpha(:j0_ffs(iky, ikx, iz, ivmu)%max_idx)
                  !                   call test_ffs_bessel_coefs (j0_ffs(iky,ikx,iz,ivmu)%fourier, aj0_alpha, iky, ikx, iz, j0_ffs_unit, ivmu)
                  if (.not. associated(j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%fourier)) &
                     allocate (j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%fourier(j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%max_idx))
                  !> fill the array with the requisite coefficients
                  j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%fourier = j0_B_maxwell_kalpha(:j0_B_maxwell_ffs(iky, ikx, iz, ivmu)%max_idx)
                  !                   call test_ffs_bessel_coefs (j0_B_maxwell_ffs(iky,ikx,iz,ivmu)%fourier, j0_B_maxwell, iky, ikx, iz, j0_B_maxwell_ffs_unit, ivmu)
               end do
            end do
         end do
      end do
      deallocate (aj0_alpha, j0_B_maxwell, j0_B_maxwell_kalpha)

      !> calculate the reduction factor of Fourier modes
      !> used to represent J0
      !> avoid overflow by converting integers to reals before multiplying
      rtmp = real(naky) * real(naky_all) * real(ikx_max) * real(nztot) * real(nmu) * real(nvpa) * real(nspec)
      call sum_allreduce(ia_max_j0_count)
      ia_max_j0_reduction_factor = real(ia_max_j0_count) / rtmp
      call sum_allreduce(ia_max_j0_B_maxwell_count)
      ia_max_j0_B_maxwell_reduction_factor = real(ia_max_j0_B_maxwell_count) / rtmp

      if (proc0) then
         write (*, *) 'average number of k-alphas needed to represent J0(kperp(alpha))=', ia_max_j0_reduction_factor * naky, 'out of ', naky
         write (*, *) 'average number of k-alphas needed to represent J0(kperp(alpha))*B(alpha)*exp(-v^2)=', &
            ia_max_j0_B_maxwell_reduction_factor * naky, 'out of ', naky
         write (*, *)
      end if

      deallocate (wgts)
      deallocate (aj0_kalpha)
      deallocate (kperp2_swap)

      !       call close_output_file (j0_ffs_unit)
      !       call close_output_file (j0_B_maxwell_ffs_unit)

   contains

      !> inverse fourier transform coefs%fourier for several phase space points and compare with
      !> unfiltered version in alpha-space
      ! subroutine test_ffs_bessel_coefs (coefs, f_alpha, iky, ikx, iz, unit, ivmu)

      !   use stella_layouts, only: vmu_lo, iv_idx, is_idx, imu_idx

      !   implicit none

      !   complex, dimension (:), intent (in) :: coefs
      !   real, dimension (:), intent (in) :: f_alpha
      !   integer, intent (in) :: iky, ikx, iz
      !   integer, intent (in) :: unit
      !   integer, intent (in), optional :: ivmu

      !   integer :: iv, imu, is

      !   if (present(ivmu)) then
      !      !> coefficients should all be independent of vpa, so only do comparison for one vpa point
      !      iv = iv_idx(vmu_lo,ivmu)
      !      if (iv == 1) then
      !         !> only sample subset of mu locations
      !         imu = imu_idx(vmu_lo,ivmu)
      !         if (mod(imu-1,nmu/2-1)==0) then
      !            is = is_idx(vmu_lo,ivmu)
      !            call test_ffs_bessel_coefs_subset (coefs, f_alpha, iky, ikx, iz, unit, iv, imu, is)
      !         end if
      !      end if
      !   else
      !      call test_ffs_bessel_coefs_subset (coefs, f_alpha, iky, ikx, iz, unit)
      !   end if

      ! end subroutine test_ffs_bessel_coefs

      ! subroutine test_ffs_bessel_coefs_subset (coefs, f_alpha, iky, ikx, iz, unit, iv, imu, is)

      !   use constants, only: pi
      !   use zgrid, only: nzgrid, zed
      !   use kt_grids, only: naky, nalpha, aky_all_ordered
      !   use vpamu_grids, only: mu
      !   use stella_transforms, only: transform_kalpha2alpha
      !   use stella_geometry, only: alpha

      !   implicit none

      !   complex, dimension (:), intent (in) :: coefs
      !   real, dimension (:), intent (in) :: f_alpha
      !   integer, intent (in) :: iky, ikx, iz
      !   integer, intent (in) :: unit
      !   integer, intent (in), optional :: iv, imu, is

      !   complex, dimension (:), allocatable :: coefs_padded
      !   real, dimension (:), allocatable :: f_alpha_approx

      !   integer :: ia
      !   integer :: max_idx
      !   real :: relative_error
      !   real, parameter :: minval = 1.0e-3

      !   ! only sample a subset of z locations
      !   if (mod(iz,nzgrid/2)==0) then
      !      ! consider only kx = 0
      !      if (ikx == 1) then
      !         allocate (coefs_padded(naky))
      !         allocate (f_alpha_approx(nalpha))
      !         ! initialize the padded coefficient array to zero
      !         coefs_padded = 0.0
      !         ! fill in non-zero entries with truncated Fourier coefficients
      !         max_idx = size(coefs)
      !         coefs_padded(:max_idx) = coefs
      !         ! inverse Fourier transform to get alpha-dependent function
      !         call transform_kalpha2alpha (coefs_padded, f_alpha_approx)
      !         if (present(iv)) then
      !            do ia = 1, nalpha
      !               relative_error = 2.0*abs(f_alpha(ia)-f_alpha_approx(ia))/(abs(f_alpha(ia)) + abs(f_alpha_approx(ia)))
      !               write (unit,*) alpha(ia), f_alpha(ia), f_alpha_approx(ia), &
      !                    relative_error, aky_all_ordered(iky), ikx, iz, zed(iz), iv, imu, mu(imu), is
      !            end do
      !            ! user 2*pi periodicity in alpha to fill in final point (for visualization purposes)
      !            ia = 1
      !            relative_error = 2.0*abs(f_alpha(ia)-f_alpha_approx(ia))/(abs(f_alpha(ia)) + abs(f_alpha_approx(ia)))
      !            write (unit,*) 2.0*pi, f_alpha(1), f_alpha_approx(1), &
      !                 relative_error, aky_all_ordered(iky), ikx, iz, zed(iz), iv, imu, mu(imu), is
      !         else
      !            do ia = 1, nalpha
      !               relative_error = 2.0*abs(f_alpha(ia)-f_alpha_approx(ia))/(abs(f_alpha(ia)) + abs(f_alpha_approx(ia)))
      !               write (unit,*) alpha(ia), f_alpha(ia), f_alpha_approx(ia), &
      !                    relative_error, aky_all_ordered(iky), ikx, iz, zed(iz)
      !            end do
      !            ! user 2*pi periodicity in alpha to fill in final point (for visualization purposes)
      !            ia = 1
      !            relative_error = 2.0*abs(f_alpha(ia)-f_alpha_approx(ia))/(abs(f_alpha(ia)) + abs(f_alpha_approx(ia)))
      !            write (unit,*) 2.0*pi, f_alpha(1), f_alpha_approx(1), &
      !                 relative_error, aky_all_ordered(iky), ikx, iz, zed(iz)
      !         end if
      !         write (unit,*)
      !         deallocate (coefs_padded, f_alpha_approx)
      !      end if
      !   end if

      ! end subroutine test_ffs_bessel_coefs_subset

   end subroutine init_bessel_ffs

   !> subroutine takes a set of Fourier coefficients (ft)
   !> and returns the minimum number of coeffients that must be retained (idx)
   !> to ensure that the relative error in the total spectral energy is
   !> below a specified tolerance (tol_floor)
   subroutine find_max_required_kalpha_index(ft, idx, imu, iz, is)

      use vpamu_grids, only: maxwell_mu

      implicit none

      complex, dimension(:), intent(in) :: ft
      integer, intent(out) :: idx
      integer, intent(in), optional :: imu, iz, is

      real, parameter :: tol_floor = 1.0e-8
      integer :: i, n
      real :: subtotal, total
      real :: tol
      real, dimension(:), allocatable :: ftmod2

      n = size(ft)

      ! use conservative estimate
      ! when deciding number of modes to retain
      if (present(imu) .and. present(iz) .and. present(is)) then
         !       tol = min(0.1,tol_floor/maxval(maxwell_mu(:,iz,imu,is)))
         tol = min(1.0e-6, tol_floor / maxval(maxwell_mu(:, iz, imu, is)))
      else
         tol = tol_floor
      end if

      allocate (ftmod2(n))
      ! get spectral energy associated with each mode
      ftmod2 = sqrt(real(ft * conjg(ft)))
      ! get total spectral energy
      total = sqrt(sum(ftmod2))
      subtotal = 0.

      ! find minimum spectral index for which
      ! desired percentage of spectral energy contained
      ! in modes with indices at or below it
      if (total > 0.) then
         i = 1
         do while (subtotal < total * (1.0 - tol))
            idx = i
            subtotal = sqrt(sum(ftmod2(:i)))
            i = i + 1
         end do
      else
         idx = 1
      end if

      deallocate (ftmod2)

   end subroutine find_max_required_kalpha_index

   subroutine finish_bessel

      implicit none

      if (allocated(aj0v)) deallocate (aj0v)
      if (allocated(aj1v)) deallocate (aj1v)
      if (allocated(aj0x)) deallocate (aj0x)
      if (allocated(aj1x)) deallocate (aj1x)
      if (allocated(j0_ffs)) deallocate (j0_ffs)
      if (allocated(j0_B_maxwell_ffs)) deallocate (j0_B_maxwell_ffs)

      bessinit = .false.

   end subroutine finish_bessel

   subroutine gyro_average_kxky_local(field, iz, ivmu, gyro_field)

      use physics_flags, only: full_flux_surface

      implicit none

      complex, dimension(:, :), intent(in) :: field
      integer, intent(in) :: iz, ivmu
      complex, dimension(:, :), intent(out) :: gyro_field

      if (full_flux_surface) then
         !> if simulating a full flux surface, the alpha dependence present
         !> in kperp makes gyro-averaging non-local in k-space
         call gyro_average(field, gyro_field, j0_ffs(:, :, iz, ivmu))
      else
         !> if simulating a flux tube, a gyro-average is local in k-space
         gyro_field = aj0x(:, :, iz, ivmu) * field
      end if

   end subroutine gyro_average_kxky_local

   subroutine gyro_average_kxkyz_local(field, ivmu, gyro_field)

      use physics_flags, only: full_flux_surface
      use zgrid, only: nzgrid, ntubes

      implicit none

      complex, dimension(:, :, -nzgrid:, :), intent(in) :: field
      integer, intent(in) :: ivmu
      complex, dimension(:, :, -nzgrid:, :), intent(out) :: gyro_field

      if (full_flux_surface) then
         !> if simulating a full flux surface, the alpha dependence present
         !> in kperp makes gyro-averaging non-local in k-space
         call gyro_average(field, gyro_field, j0_ffs(:, :, :, ivmu))
      else
         !> if simulating a flux tube, a gyro-average is local in k-space
         gyro_field = spread(aj0x(:, :, :, ivmu), 4, ntubes) * field
      end if

   end subroutine gyro_average_kxkyz_local

   subroutine gyro_average_ffs_kxky_local(field, gyro_field, coefs)

      use kt_grids, only: naky, nakx
      use kt_grids, only: naky_all, ikx_max
      use kt_grids, only: swap_kxky_ordered, swap_kxky_back_ordered

      implicit none

      complex, dimension(:, :), intent(in) :: field
      complex, dimension(:, :), intent(out) :: gyro_field
      type(coupled_alpha_type), dimension(:, :), intent(in) :: coefs

      integer :: iky, ikx, ikyp
      integer :: idx
      complex, dimension(:, :), allocatable :: field_kyall, gyro_field_kyall

      ! need to switch from ky>=0 and all kx
      ! to kx>=0 and all ky (using reality condition)
      allocate (field_kyall(naky_all, ikx_max))
      allocate (gyro_field_kyall(naky_all, ikx_max)); gyro_field_kyall = 0.
      call swap_kxky_ordered(field, field_kyall)
      ! NB: J0(kx,ky) = J0(-kx,-ky) and Gamma0(kx,ky) = Gamma0(-kx,-ky)
      do ikx = 1, ikx_max
         do iky = 1, naky_all
            ! account for contributions from less positive ky values (and this ky itself)
            do ikyp = 1, min(naky, iky)
               ! idx is the index corresponding to k_alpha - k_alpha'
               ! runs from iky down to 1
               idx = iky - ikyp + 1
               ! if the Fourier coefficient corresponding to this value of (k_alpha-k_alpha',k_alpha')
               ! is sufficiently small, then it will not have been included in the truncated version
               ! of the coefficients; in this case, it makes no contribution to the gyro-average sum
               if (coefs(idx, ikx)%max_idx >= ikyp) then
                  gyro_field_kyall(iky, ikx) = gyro_field_kyall(iky, ikx) &
                                               + coefs(idx, ikx)%fourier(ikyp) * field_kyall(idx, ikx)
               end if
            end do
            ! if iky = naky_all, then already at max positive ky, so no contributions
            ! from more positive ky value possible
            if (iky == naky_all) cycle
            ! account for contributions from more positive ky values (but not this ky itself,
            ! as already accounted for above
            do ikyp = 2, min(naky, naky_all - iky + 1)
               ! idx is the index corresponding to k_alpha - k_alpha'
               ! runs from iky + 1 up to iky + naky (or until naky_all, if it is reached first)
               idx = iky + ikyp - 1
               ! if the Fourier coefficient corresponding to this value of (k_alpha-k_alpha',k_alpha')
               ! is sufficiently small, then it will not have been included in the truncated version
               ! of the coefficients; in this case, it makes no contribution to the gyro-average sum
               if (coefs(idx, ikx)%max_idx >= ikyp) then
                  ! the k_alpha' values considered in this loop are negative, but only have
                  ! Fourier coefficients for positive ky values;
                  ! must use the reality condition to convert this to the equivalent coefficients for negative ky
                  gyro_field_kyall(iky, ikx) = gyro_field_kyall(iky, ikx) &
                                               + conjg(coefs(idx, ikx)%fourier(ikyp)) * field_kyall(idx, ikx)
               end if
            end do
         end do
      end do

      call swap_kxky_back_ordered(gyro_field_kyall, gyro_field)
      deallocate (field_kyall, gyro_field_kyall)

   end subroutine gyro_average_ffs_kxky_local

   subroutine gyro_average_ffs_kxkyz_local(field, gyro_field, coefs)

      use zgrid, only: nzgrid, ntubes

      implicit none

      complex, dimension(:, :, -nzgrid:, :), intent(in) :: field
      complex, dimension(:, :, -nzgrid:, :), intent(out) :: gyro_field
      type(coupled_alpha_type), dimension(:, :, -nzgrid:), intent(in) :: coefs

      integer :: iz, it

      do it = 1, ntubes
         do iz = -nzgrid, nzgrid
            call gyro_average(field(:, :, iz, it), gyro_field(:, :, iz, it), coefs(:, :, iz))
         end do
      end do

   end subroutine gyro_average_ffs_kxkyz_local

   subroutine gyro_average_ffs(dist, gyro_dist, coefs)

      use common_types, only: coupled_alpha_type
      use stella_layouts, only: vmu_lo
      use zgrid, only: nzgrid

      implicit none

      complex, dimension(:, :, -nzgrid:, :, vmu_lo%llim_proc:), intent(in) :: dist
      complex, dimension(:, :, -nzgrid:, :, vmu_lo%llim_proc:), intent(out) :: gyro_dist
      type(coupled_alpha_type), dimension(:, :, -nzgrid:, vmu_lo%llim_proc:), intent(in) :: coefs

      integer :: ivmu

      do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
         call gyro_average(dist(:, :, :, :, ivmu), gyro_dist(:, :, :, :, ivmu), coefs(:, :, :, ivmu))
      end do

   end subroutine gyro_average_ffs

   subroutine gyro_average_vmu_local(distfn, ikxkyz, gyro_distfn)

      use vpamu_grids, only: nvpa

      implicit none

      complex, dimension(:, :), intent(in) :: distfn
      integer, intent(in) :: ikxkyz
      complex, dimension(:, :), intent(out) :: gyro_distfn

      gyro_distfn = spread(aj0v(:, ikxkyz), 1, nvpa) * distfn

   end subroutine gyro_average_vmu_local

   subroutine gyro_average_vmus_nonlocal(field, iky, ikx, iz, gyro_field)

      use stella_layouts, only: vmu_lo

      implicit none

      complex, dimension(vmu_lo%llim_proc:), intent(in) :: field
      integer, intent(in) :: iky, ikx, iz
      complex, dimension(vmu_lo%llim_proc:), intent(out) :: gyro_field

      gyro_field = aj0x(iky, ikx, iz, :) * field

   end subroutine gyro_average_vmus_nonlocal

   subroutine gyro_average_j1_kxky_local(field, iz, ivmu, gyro_field)

      implicit none

      complex, dimension(:, :), intent(in) :: field
      integer, intent(in) :: iz, ivmu
      complex, dimension(:, :), intent(out) :: gyro_field

      gyro_field = aj1x(:, :, iz, ivmu) * field

   end subroutine gyro_average_j1_kxky_local

   subroutine gyro_average_j1_kxkyz_local(field, ivmu, gyro_field)

      use zgrid, only: nzgrid, ntubes

      implicit none

      complex, dimension(:, :, -nzgrid:, :), intent(in) :: field
      integer, intent(in) :: ivmu
      complex, dimension(:, :, -nzgrid:, :), intent(out) :: gyro_field

      integer :: iz, it

      do it = 1, ntubes
         do iz = -nzgrid, nzgrid
            gyro_field(:, :, iz, it) = aj1x(:, :, iz, ivmu) * field(:, :, iz, it)
         end do
      end do

   end subroutine gyro_average_j1_kxkyz_local

   subroutine gyro_average_j1_vmu_local(distfn, ikxkyz, gyro_distfn)

      use vpamu_grids, only: nvpa

      implicit none

      complex, dimension(:, :), intent(in) :: distfn
      integer, intent(in) :: ikxkyz
      complex, dimension(:, :), intent(out) :: gyro_distfn

      gyro_distfn = spread(aj1v(:, ikxkyz), 1, nvpa) * distfn

   end subroutine gyro_average_j1_vmu_local

   subroutine band_lu_solve_ffs(lu, solvec)

      use common_types, only: gam0_ffs_type
      use zgrid, only: nzgrid
      use kt_grids, only: ikx_max

      implicit none

      type(gam0_ffs_type), dimension(:, -nzgrid:), intent(in) :: lu
      complex, dimension(:, :, -nzgrid:), intent(in out) :: solvec

      integer :: ikx, iz

      do iz = -nzgrid, nzgrid
         do ikx = 1, ikx_max
            call band_lu_solve_ffs_single(lu(ikx, iz), solvec(:, ikx, iz))
         end do
      end do

   end subroutine band_lu_solve_ffs

   subroutine band_lu_solve_ffs_single(lu, solvec)

      use common_types, only: gam0_ffs_type
      use kt_grids, only: naky

      implicit none

      type(gam0_ffs_type), intent(in) :: lu
      complex, dimension(:), intent(in out) :: solvec

      integer :: n, nsubdiag, nsupdiag, nrhs
      integer :: info
      complex, dimension(:, :), allocatable :: solmat

      ! n is the order of the matrix for which we have the LU factorisation
      n = size(lu%pivot_index)
      ! nsubdiag and nsupdiag are the number of sub- and super-diagonals within the band of the matrix to be decomposed
      nsubdiag = naky - 1; nsupdiag = naky - 1
      ! nrhs is the number of right-hand sides of the matrix equation lu%matrix * solvec = rhs for which to solve
      nrhs = 1
      ! initialise solmat = rhs, as it will be overwritten by zgbtrs below
      allocate (solmat(size(solvec), 1))
      solmat(:, 1) = solvec

      call zgbtrs('N', n, nsubdiag, nsupdiag, nrhs, lu%matrix, size(lu%matrix, 1), lu%pivot_index, solmat, size(solmat, 1), info)

      solvec = solmat(:, 1)

      deallocate (solmat)

   end subroutine band_lu_solve_ffs_single

   subroutine band_lu_factorisation_ffs(gam0, lu_gam0)

      use common_types, only: coupled_alpha_type, gam0_ffs_type
      use zgrid, only: nzgrid
      use kt_grids, only: ikx_max, naky_all, naky

      implicit none

      type(coupled_alpha_type), dimension(:, :, -nzgrid:), intent(in) :: gam0
      type(gam0_ffs_type), dimension(:, -nzgrid:), intent(out) :: lu_gam0

      integer :: iky, ikx, iz

      complex, dimension(:, :), allocatable :: gam0tmp

      allocate (gam0tmp(naky, naky_all))

      do ikx = 1, ikx_max
         do iz = -nzgrid, nzgrid
            ! create array from Fourier coefficients of Gamma_0(ky,y)
            do iky = 1, naky_all
               gam0tmp(:, iky) = gam0(iky, ikx, iz)%fourier
            end do
            call band_lu_factorisation_single(gam0tmp, lu_gam0(ikx, iz))
         end do
      end do

      deallocate (gam0tmp)

   end subroutine band_lu_factorisation_ffs

   subroutine band_lu_factorisation_single(gam0, lu_gam0)

      use common_types, only: gam0_ffs_type
      use kt_grids, only: naky, naky_all

      implicit none

      complex, dimension(:, :), intent(in) :: gam0
      type(gam0_ffs_type), intent(out) :: lu_gam0

      integer :: nrows, ncols, nsubdiag, nsupdiag, leading_dim
      integer :: info
      integer :: i, imod

      ! nrows and ncols are the number of rows and columns of the matrix to be LU-decomposed (variant of gam0)
      ! this matrix is naky_all x naky_all
      nrows = naky_all; ncols = naky_all
      ! nsubdiag and nsupdiag are the number of sub- and super-diagonals within the band of the matrix to be decomposed
      nsubdiag = naky - 1; nsupdiag = naky - 1
      ! leading_dim is the 'leading dimension' of the lu_gam0 array
      leading_dim = 2 * nsubdiag + nsupdiag + 1

      ! lu_gam0 is a re-arranged version of gam0 on entry, and on exit contains details of LU factorisation
      ! that can be used by zgbtrs to solve the linear system gam0 * phi = rhs
      if (.not. associated(lu_gam0%matrix)) then
         allocate (lu_gam0%matrix(leading_dim, ncols))
         ! initialise first nsubdiag rows to zero, as they are unused
         lu_gam0%matrix(:nsubdiag, :) = 0.0
         ! fill next supdiag rows using elements from super-diagonals
         ! using reality of gam0 to set fourier(ky < 0) = conjg(fourier(ky > 0))
         do i = 1, nsupdiag
            imod = naky - i + 1
            lu_gam0%matrix(nsubdiag + i, imod:) = conjg(gam0(imod, imod:))
            ! fill unused entries with zero
            lu_gam0%matrix(nsubdiag + i, :imod - 1) = 0.0
         end do
         ! fill next row using main diagonal entries
         lu_gam0%matrix(nsubdiag + nsupdiag + 1, :) = gam0(1, :)
         ! fill remaining nsubdiag rows using elements from sub-diagonals
         do i = 1, nsubdiag
            imod = naky + i - 1
            lu_gam0%matrix(leading_dim - i + 1, :imod) = gam0(naky - i + 1, :imod)
            ! fill unused entries with zeroes
            lu_gam0%matrix(leading_dim - i + 1, imod + 1:) = 0.0
         end do
      end if

      if (.not. associated(lu_gam0%pivot_index)) allocate (lu_gam0%pivot_index(min(nrows, ncols)))
      ! overwrites lu_gam0%matrix with information needed to solve the linear system gam0 * phi = rhs;
      ! also returns pivot_index, with pivot_index(i) giving the row of the matrix interchanged with the ith row,
      ! and info, which should be zero if the LU factorisation is successful
      call zgbtrf(nrows, ncols, nsubdiag, nsupdiag, lu_gam0%matrix, leading_dim, lu_gam0%pivot_index, info)

   end subroutine band_lu_factorisation_single

   ! subroutine test_band_lu_factorisation (gam0, lu_gam0)

   !   use common_types, only: coupled_alpha_type, gam0_ffs_type
   !   use zgrid, only: nzgrid
   !   use kt_grids, only: naky_all, naky

   !   implicit none

   !   type (coupled_alpha_type), dimension (:,:,-nzgrid:), intent (in) :: gam0
   !   type (gam0_ffs_type), dimension (:,-nzgrid:), intent (out) :: lu_gam0

   !   integer :: iky, ikx, ikyp, iz
   !   complex, dimension (naky_all) :: solvec

   !   ikx = 1 ; iz = -nzgrid
   !   do iky = 1, naky_all
   !      do ikyp = 1, naky
   !         gam0(iky,ikx,iz)%fourier(ikyp) = iky-naky + ikyp-1
   !      end do
   !   end do

   !   call band_lu_factorisation_ffs (gam0, lu_gam0)

   !   do iky = 1, naky_all
   !      solvec(iky) = iky
   !   end do

   !   call band_lu_solve_ffs_single (lu_gam0(ikx,iz), solvec)

   !   do iky = 1, naky_all
   !      write (*,*) 'iky: ', iky, 'solution: ', solvec(iky)
   !   end do
   !   stop

   ! end subroutine test_band_lu_factorisation

end module gyro_averages