module stella_geometry

  use common_types, only: flux_surface_type

  implicit none

  public :: init_geometry, finish_init_geometry, finish_geometry
  public :: communicate_geo_multibox
  public :: grho
  public :: bmag, dbdzed, btor, bmag_psi0
  public :: gradpar, gradpar_eqarc, zed_eqarc
  public :: cvdrift, cvdrift0
  public :: gbdrift, gbdrift0
  public :: dcvdriftdrho, dcvdrift0drho
  public :: dgbdriftdrho, dgbdrift0drho
  public :: gds2, gds21, gds22, gds23, gds24, gds25, gds26
  public :: dgds2dr, dgds21dr, dgds22dr
  public :: exb_nonlin_fac, exb_nonlin_fac_p
  public :: jacob, djacdrho
  public :: drhodpsi, drhodpsi_psi0
  public :: dl_over_b, d_dl_over_b_drho
  public :: dBdrho, d2Bdrdth, dgradpardrho, dIdrho
  public :: geo_surf
  public :: Rmajor
  public :: alpha
  public :: theta_vmec
  public :: zeta
  public :: zed_scalefac
  public :: dxdXcoord, dydalpha
  public :: aref, bref
  public :: twist_and_shift_geo_fac
  public :: q_as_x, get_x_to_rho, gfac

  private

  type (flux_surface_type) :: geo_surf

  real :: aref, bref
  real :: dxdXcoord, dydalpha
  real :: dqdrho
  real :: dIdrho
  real :: drhodpsi, drhodpsi_psi0, shat, qinp
  real :: exb_nonlin_fac, exb_nonlin_fac_p
  real :: gradpar_eqarc
  real :: zed_scalefac
  real :: twist_and_shift_geo_fac, gfac
  real, dimension (:), allocatable :: zed_eqarc
  real, dimension (:), allocatable :: gradpar
  real, dimension (:,:), allocatable :: bmag, bmag_psi0, dbdzed
  real, dimension (:,:), allocatable :: cvdrift, cvdrift0
  real, dimension (:,:), allocatable :: gbdrift, gbdrift0
  real, dimension (:,:), allocatable :: dcvdriftdrho, dcvdrift0drho
  real, dimension (:,:), allocatable :: dgbdriftdrho, dgbdrift0drho
  real, dimension (:,:), allocatable :: gds2, gds21, gds22, gds23, gds24, gds25, gds26
  real, dimension (:,:), allocatable :: dgds2dr, dgds21dr
  real, dimension (:,:), allocatable :: dgds22dr
  real, dimension (:,:), allocatable :: theta_vmec
  real, dimension (:,:), allocatable :: jacob, djacdrho, grho
  real, dimension (:,:), allocatable :: dl_over_b, d_dl_over_b_drho
  real, dimension (:), allocatable :: dBdrho, d2Bdrdth, dgradpardrho
  real, dimension (:), allocatable :: btor, Rmajor
  real, dimension (:), allocatable :: alpha
  real, dimension (:,:), allocatable :: zeta

  integer :: geo_option_switch
  integer, parameter :: geo_option_local = 1
  integer, parameter :: geo_option_inputprof = 2
  integer, parameter :: geo_option_vmec = 3
  integer, parameter :: geo_option_multibox = 4

  logical :: overwrite_geometry
  logical :: overwrite_bmag, overwrite_gradpar
  logical :: overwrite_gds2, overwrite_gds21, overwrite_gds22
  logical :: overwrite_gds23, overwrite_gds24
  logical :: overwrite_gbdrift, overwrite_cvdrift, overwrite_gbdrift0
  logical :: q_as_x
  character (100) :: geo_file

  logical :: geoinit = .false.

contains

  subroutine init_geometry (nalpha)

    use constants, only: pi
    use mp, only: proc0
    use millerlocal, only: read_local_parameters, get_local_geo
    use millerlocal, only: communicate_parameters_multibox
    use vmec_geo, only: read_vmec_parameters, get_vmec_geo
    use inputprofiles_interface, only: read_inputprof_geo
    use zgrid, only: nzed, nzgrid
    use zgrid, only: zed, delzed
    use zgrid, only: shat_zero, zed_equal_arc
    use zgrid, only: boundary_option_switch, boundary_option_self_periodic
    use file_utils, only: get_unused_unit
    use physics_flags, only: include_geometric_variation

    implicit none

    logical, parameter :: debug = .false.

    integer, intent (in) :: nalpha

    real :: dpsidrho, dpsidrho_psi0
    integer :: iy
    integer :: sign_torflux
    integer :: dxdXcoord_sign, dydalpha_sign
    real :: field_period_ratio
    real, dimension (:,:), allocatable :: grad_alpha_grad_alpha
    real, dimension (:,:), allocatable :: grad_alpha_grad_psi
    real, dimension (:,:), allocatable :: grad_psi_grad_psi
    real, dimension (:,:), allocatable :: gbdrift_alpha, cvdrift_alpha
    real, dimension (:,:), allocatable :: gbdrift0_psi, cvdrift0_psi

    if (geoinit) return
    geoinit = .true.

    ! B = grad alpha x grad psi
    ! for tokamak calculations, alpha = zeta - q * theta
    ! and psi = psi_poloidal
    ! for stellarator calculations, alpha = theta - iota * zeta
    ! and psi = -psi_toroidal

    ! default is no re-scaling of zed
    zed_scalefac = 1.0

    if (proc0) then
       call read_parameters
       select case (geo_option_switch)
       case (geo_option_local)
          ! read in Miller local parameters
          call read_local_parameters (nzed,nzgrid,geo_surf)
          ! allocate geometry arrays
          call allocate_arrays (nalpha, nzgrid)
          ! use Miller local parameters to get
          ! geometric coefficients needed by stella
          call get_local_geo (nzed, nzgrid, zed, zed_equal_arc, &
               dpsidrho, dpsidrho_psi0, dIdrho, grho(1,:), &
               bmag(1,:), bmag_psi0(1,:), &
               gds2(1,:), gds21(1,:), gds22(1,:), &
               gds23(1,:), gds24(1,:), gradpar, &
               gbdrift0(1,:), gbdrift(1,:), cvdrift0(1,:), cvdrift(1,:), &
               dBdrho, d2Bdrdth, dgradpardrho, btor, rmajor, &
               dcvdrift0drho(1,:), dcvdriftdrho(1,:), &
               dgbdrift0drho(1,:), dgbdriftdrho(1,:), &
               dgds2dr(1,:),dgds21dr(1,:), &
               dgds22dr(1,:), djacdrho(1,:))
          ! note that psi here is the enclosed poloidal flux divided by 2pi
          drhodpsi = 1./dpsidrho
          drhodpsi_psi0 = 1./dpsidrho_psi0
          ! dxdXcoord = a*Bref*dx/dpsi = sign(dx/dpsi) * a*q/r
          dxdXcoord_sign = 1
          if(q_as_x) then
            dxdXcoord = dxdXcoord_sign*dpsidrho
          else
            dxdXcoord = dxdXcoord_sign*geo_surf%qinp_psi0/geo_surf%rhoc_psi0
          endif
          ! dydalpha = (dy/dalpha) / a = sign(dydalpha) * (dpsi/dr) / (a*Bref)
          dydalpha_sign = 1
          dydalpha = dydalpha_sign*dpsidrho
          ! abs(twist_and_shift_geo_fac) is dkx/dky * jtwist
          ! minus its sign gives the direction of the shift in kx
          ! to be used for twist-and-shift BC
          if(q_as_x) then
            twist_and_shift_geo_fac = 2.0*pi
          else
            twist_and_shift_geo_fac = 2.0*pi*geo_surf%shat_psi0
          endif
          ! aref and bref should not be needed, so set to 1
          aref = 1.0 ; bref = 1.0
          zeta(1,:) = zed*geo_surf%qinp
       case (geo_option_multibox)
          ! read in Miller local parameters
          call read_local_parameters (nzed,nzgrid,geo_surf)

          ! allocate geometry arrays
          call allocate_arrays (nalpha, nzgrid)

          call communicate_parameters_multibox(surf=geo_surf)

          ! use Miller local parameters to get
          ! geometric coefficients needed by stella
          call get_local_geo (nzed, nzgrid, zed, zed_equal_arc, &
               dpsidrho, dpsidrho_psi0, dIdrho, grho(1,:), &
               bmag(1,:), bmag_psi0(1,:), &
               gds2(1,:), gds21(1,:), gds22(1,:), &
               gds23(1,:), gds24(1,:), gradpar, &
               gbdrift0(1,:), gbdrift(1,:), cvdrift0(1,:), cvdrift(1,:), &
               dBdrho, d2Bdrdth, dgradpardrho, btor, rmajor, &
               dcvdrift0drho(1,:), dcvdriftdrho(1,:), &
               dgbdrift0drho(1,:), dgbdriftdrho(1,:), &
               dgds2dr(1,:),dgds21dr(1,:), &
               dgds22dr(1,:), djacdrho(1,:))
          ! note that psi here is the enclosed poloidal flux divided by 2pi
          drhodpsi = 1./dpsidrho
          drhodpsi_psi0 = 1./dpsidrho_psi0
          ! dxdXcoord = a*Bref*dx/dpsi = sign(dx/dpsi) * a*q/r
          dxdXcoord_sign = 1
          if(q_as_x) then
            dxdXcoord = dxdXcoord_sign/drhodpsi_psi0
          else
            dxdXcoord = dxdXcoord_sign*geo_surf%qinp_psi0/geo_surf%rhoc_psi0
          endif
          ! dydalpha = (dy/dalpha) / a = sign(dydalpha) * (dpsi/dr) / (a*Bref)
          dydalpha_sign = 1
          dydalpha = dydalpha_sign/drhodpsi_psi0
          ! abs(twist_and_shift_geo_fac) is dkx/dky * jtwist
          ! minus its sign gives the direction of the shift in kx
          ! to be used for twist-and-shift BC
          if(q_as_x) then
            twist_and_shift_geo_fac = 2.0*pi
          else
            twist_and_shift_geo_fac = 2.0*pi*geo_surf%shat_psi0
          endif
          ! aref and bref should not be needed, so set to 1
          aref = 1.0 ; bref = 1.0
          zeta(1,:) = zed*geo_surf%qinp

       case (geo_option_inputprof)
          ! first read in some local parameters
          ! only thing needed really is rhoc
          call read_local_parameters (nzed,nzgrid,geo_surf)
          ! allocate geometry arrays
          call allocate_arrays (nalpha, nzgrid)
          ! now overwrite local parameters
          ! with those from input.profiles file
          ! use rhoc from input as surface
          call read_inputprof_geo (geo_surf)
          call get_local_geo (nzed, nzgrid, zed, zed_equal_arc, &
               dpsidrho, dpsidrho_psi0, dIdrho, grho(1,:), &
               bmag(1,:), bmag_psi0(1,:), &
               gds2(1,:), gds21(1,:), gds22(1,:), &
               gds23(1,:), gds24(1,:), gradpar, &
               gbdrift0(1,:), gbdrift(1,:), cvdrift0(1,:), cvdrift(1,:), &
               dBdrho, d2Bdrdth, dgradpardrho, btor, rmajor, &
               dcvdrift0drho(1,:), dcvdriftdrho(1,:), &
               dgbdrift0drho(1,:), dgbdriftdrho(1,:), &
               dgds2dr(1,:),dgds21dr(1,:), &
               dgds22dr(1,:),djacdrho(1,:))
          ! psi here is enclosed poloidal flux divided by 2pi
          drhodpsi = 1./dpsidrho
          drhodpsi_psi0 = 1./dpsidrho_psi0
          ! dxdXcoord = a*Bref*dx/dpsi = sign(dx/dpsi) * a*q/r
          dxdXcoord_sign = 1
          dxdXcoord = dxdXcoord_sign*geo_surf%qinp/geo_surf%rhoc
          ! dydalpha = (dy/dalpha) / a = sign(dydalpha) * (dpsi/dr) / (a*Bref)
          dydalpha_sign = 1
          dydalpha = dydalpha_sign*dpsidrho
          ! abs(twist_and_shift_geo_fac) is dkx/dky * jtwist
          ! minus its sign gives the direction of the shift in kx
          ! to be used for twist-and-shift BC
          twist_and_shift_geo_fac = 2.0*pi*geo_surf%shat
          ! aref and bref should not be needed so set to 1
          aref = 1.0 ; bref = 1.0

          zeta(1,:) = zed*geo_surf%qinp

       case (geo_option_vmec)
          ! read in input parameters for vmec
          ! nalpha may be specified via input file
          if (debug) write (*,*) 'init_geometry::read_vmec_parameters'
          call read_vmec_parameters
          ! allocate geometry arrays
          if (debug) write (*,*) 'init_geometry::allocate_arrays'
          call allocate_arrays (nalpha, nzgrid)
          if (debug) write (*,*) 'init_geometry::allocate_temporary_arrays'
          call allocate_temporary_arrays (nalpha, nzgrid)
          ! get geometry coefficients from vmec
          if (debug) write (*,*) 'init_geometry::get_vmec_geo'
          call get_vmec_geo (nzgrid, nalpha, geo_surf, grho, bmag, gradpar, grad_alpha_grad_alpha, &
               grad_alpha_grad_psi, grad_psi_grad_psi, &
               gds23, gds24, gds25, gds26, gbdrift_alpha, gbdrift0_psi, &
               cvdrift_alpha, cvdrift0_psi, sign_torflux, &
               theta_vmec, zed_scalefac, aref, bref, alpha, zeta, &
               field_period_ratio)
          ! Bref = 2*abs(psi_tor_LCFS)/a^2
          ! a*Bref*dx/dpsi_tor = sign(psi_tor)/rhotor
          ! psi = -psi_tor
          ! dxdXcoord = a*Bref*dx/dpsi = -a*Bref*dx/dpsi_tor = -sign(psi_tor)/rhotor
          dxdXcoord_sign = -1
          dxdXcoord = dxdXcoord_sign*sign_torflux/geo_surf%rhotor
          ! dydalpha = (dy/dalpha) / a = sign(dydalpha) * rhotor
          dydalpha_sign = 1
          dydalpha = dydalpha_sign*geo_surf%rhotor
          ! if using vmec, rho = sqrt(psitor/psitor_lcfs)
          ! psiN = -psitor/(aref**2*Bref)
          ! so drho/dpsiN = -drho/d(rho**2) * (aref**2*Bref/psitor_lcfs) = -1.0/rho
          drhodpsi = dxdXcoord_sign*sign_torflux/geo_surf%rhotor
          drhodpsi_psi0 = drhodpsi
          bmag_psi0 = bmag

          ! abs(twist_and_shift_geo_fac) is dkx/dky * jtwist
          ! minus its sign gives the direction of the shift in kx
          ! to be used for twist-and-shift BC
!          twist_and_shift_geo_fac = -2.*pi*geo_surf%shat*geo_surf%qinp*drhodpsi*dydalpha/(dxdpsi*geo_surf%rhotor)
          twist_and_shift_geo_fac = -2.*pi*geo_surf%shat*drhodpsi*dydalpha/(geo_surf%qinp*dxdXcoord*geo_surf%rhotor) &
               * field_period_ratio

          ! gds2 = |grad y|^2 = |grad alpha|^2 * (dy/dalpha)^2
          ! note that rhotor = sqrt(psi/psi_LCFS)
          gds2 = grad_alpha_grad_alpha * dydalpha**2
          ! gds21 = shat * grad x . grad y = shat * dx/dpsi_t * dy/dalpha * grad alpha . grad psi_t
          ! NB: psi = -psi_t and so dx/dpsi = = dx/dpsi_t, which is why there is a minus sign here
          gds21 = -grad_alpha_grad_psi * geo_surf%shat * dxdXcoord * dydalpha
          ! gds22 = shat^2 * |grad x|^2 = shat^2 * |grad psi_t|^2 * (dx/dpsi_t)^2
          gds22 = geo_surf%shat**2 * grad_psi_grad_psi * dxdXcoord**2

          ! gbdrift_alpha and cvdrift_alpha contain
          ! the grad-B and curvature drifts projected onto
          ! the grad alpha direction
          ! need the projections on grad y
          gbdrift = gbdrift_alpha * dydalpha
          cvdrift = cvdrift_alpha * dydalpha

          ! gbdrift0_psi and cvdrift0_psi contain
          ! the grad-B and curvature drifts projected onto
          ! the grad psi direction
          ! need the projections on grad x
          gbdrift0 = gbdrift0_psi * dxdXcoord
          cvdrift0 = cvdrift0_psi * dxdXcoord

          call deallocate_temporary_arrays
       end select

       if (overwrite_geometry) call overwrite_selected_geometric_coefficients (nalpha)

       ! exb_nonlin_fac is equivalent to kxfac/2 in gs2
       if(q_as_x) then
         dqdrho = geo_surf%shat * geo_surf%qinp / geo_surf%rhoc
         exb_nonlin_fac   = 0.5*dxdXcoord*dydalpha*drhodpsi*dqdrho
         !the following will get multiplied by exb_nonlin_fac in advance_exb_nonlinearity
         exb_nonlin_fac_p = geo_surf%d2qdr2/dqdrho - geo_surf%d2psidr2*drhodpsi
       else
         exb_nonlin_fac   = 0.5*dxdXcoord*dydalpha
         exb_nonlin_fac_p = 0.0
       endif
    end if

    if (.not.proc0) call allocate_arrays (nalpha, nzgrid)

    if (debug .and. proc0) write (*,*) 'init_geometry::broadcast_arrays'
    call broadcast_arrays

    gfac = 1.0
    if(.not.include_geometric_variation) gfac = 0.0

    ! should reduce to 2*pi*shat in axisymmetric case
    ! but not in non-axisymmetric case
!    twist_and_shift_geo_fac = geo_surf%shat*(gds21(1,-nzgrid)/gds22(1,-nzgrid)-gds21(1,nzgrid)/gds22(1,nzgrid))

    ! FLAG DSO - the followiing assumes a linear relation from q to rho, but this will
    !            not be correct if d2qdrho != 0
    dqdrho = geo_surf%shat * geo_surf%qinp / geo_surf%rhoc

    jacob = 1.0/abs(drhodpsi*spread(gradpar,1,nalpha)*bmag)

    ! this is dl/B
    dl_over_b = spread(delzed,1,nalpha)*jacob

    ! the next line is to avoid double counting the end points for ky = 0 modes (which leads to destabilization
    ! of the zonal modes for certain input parameters)
    ! FLAG DSO - while this is correct for ky = 0 modes and sufficient for output, if dl_over_b is applied to
    ! non-zero ky modes, a more sophisticated approach will be required that takes into account the sign of
    ! v_parallel
    dl_over_b(:,nzgrid) = 0.

    ! this is the correction to flux-surface-averaging for adiabatic electrons
    d_dl_over_b_drho = spread(delzed,1,nalpha)*djacdrho
    d_dl_over_b_drho(:,nzgrid) = 0
    d_dl_over_b_drho = d_dl_over_b_drho - dl_over_b &
                     * spread(sum(d_dl_over_b_drho,dim=2)/sum(dl_over_b,dim=2),2,2*nzgrid+1)
    d_dl_over_b_drho = gfac*d_dl_over_b_drho / spread(sum(dl_over_b,dim=2),2,2*nzgrid+1)

    ! normalize dl/B by int dl/B
    dl_over_b = dl_over_b / spread(sum(dl_over_b,dim=2),2,2*nzgrid+1)

    ! would probably be better to compute this in the various
    ! geometry subroutine (Miller, vmec, etc.), as there
    ! B is likely calculated on a finer z-grid
    do iy = 1, nalpha
       call get_dzed (nzgrid, delzed, bmag(iy,:), dbdzed(iy,:))
    end do

    ! if magnetic shear almost zero, override parallel
    ! boundary condition so that it is periodic
    if(abs(geo_surf%shat) <=  shat_zero) &
         boundary_option_switch = boundary_option_self_periodic

    ! theta_eqarc is parallel coordinate such that
    ! b . grad theta_eqarc = constant
    ! and theta_eqarc = theta at +/- pi
    ! b . grad theta_eqarc = b . grad theta dtheta_eqarc/dtheta
    ! --> dtheta_eqarc/dtheta = b . grad theta_eqarc / b . grad theta
    ! --> 2*pi = b . grad theta_eqarc * int_0^{2pi} dtheta 1/(b.grad theta)
    ! this gives b . grad theta_eqarc, from which we get
    ! theta_eqarc = theta_min + int_{0}^{theta} dtheta' b . grad theta_eqarc / b . grad theta'
    call get_gradpar_eqarc (gradpar, zed, delzed, gradpar_eqarc)
    call get_zed_eqarc (gradpar, delzed, zed, gradpar_eqarc, zed_eqarc)

    if (proc0) call write_geometric_coefficients (nalpha)

  contains

    subroutine allocate_temporary_arrays (nalpha, nzgrid)

      implicit none

      integer, intent (in) :: nalpha, nzgrid

      allocate (grad_alpha_grad_alpha(nalpha, -nzgrid:nzgrid))
      allocate (grad_alpha_grad_psi(nalpha, -nzgrid:nzgrid))
      allocate (grad_psi_grad_psi(nalpha, -nzgrid:nzgrid))
      allocate (gbdrift_alpha(nalpha, -nzgrid:nzgrid))
      allocate (cvdrift_alpha(nalpha, -nzgrid:nzgrid))
      allocate (gbdrift0_psi(nalpha, -nzgrid:nzgrid))
      allocate (cvdrift0_psi(nalpha, -nzgrid:nzgrid))

    end subroutine allocate_temporary_arrays

    subroutine deallocate_temporary_arrays

      implicit none

      deallocate (grad_alpha_grad_alpha)
      deallocate (grad_alpha_grad_psi)
      deallocate (grad_psi_grad_psi)
      deallocate (gbdrift_alpha)
      deallocate (cvdrift_alpha)
      deallocate (gbdrift0_psi)
      deallocate (cvdrift0_psi)

    end subroutine deallocate_temporary_arrays

    subroutine overwrite_selected_geometric_coefficients (nalpha)

      use file_utils, only: get_unused_unit
      use zgrid, only: nzgrid

      implicit none

      integer, intent (in) :: nalpha
      integer :: geofile_unit
      character (100) :: dum_char
      real :: dum_real

      integer :: ia, iz
      real :: bmag_file, gradpar_file
      real :: gds2_file, gds21_file, gds22_file, gds23_file, gds24_file
      real :: gbdrift_file, cvdrift_file, gbdrift0_file

      call get_unused_unit (geofile_unit)
      open (geofile_unit,file=trim(geo_file),status='old',action='read')

      read (geofile_unit,fmt=*) dum_char
      read (geofile_unit,fmt=*) dum_char
      read (geofile_unit,fmt=*) dum_char

      ! overwrite bmag, gradpar, gds2, gds21, gds22, gds23, gds24, gbdrift, cvdrift, gbdrift0, and cvdrift0
      ! with values from file
      do ia = 1, nalpha
         do iz = -nzgrid, nzgrid
            read (geofile_unit,fmt='(13e12.4)') dum_real, dum_real, dum_real, bmag_file, gradpar_file, &
                 gds2_file, gds21_file, gds22_file, gds23_file, &
                 gds24_file, gbdrift_file, cvdrift_file, gbdrift0_file
            if (overwrite_bmag) bmag(ia,iz) = bmag_file
            if (overwrite_gradpar) gradpar(iz) = gradpar_file
            if (overwrite_gds2) gds2(ia,iz) = gds2_file
            if (overwrite_gds21) gds21(ia,iz) = gds21_file
            if (overwrite_gds22) gds22(ia,iz) = gds22_file
            if (overwrite_gds23) gds23(ia,iz) = gds23_file
            if (overwrite_gds24) gds24(ia,iz) = gds24_file
            if (overwrite_gbdrift) gbdrift(ia,iz) = gbdrift_file
            if (overwrite_cvdrift) cvdrift(ia,iz) = cvdrift_file
            if (overwrite_gbdrift0) gbdrift0(ia,iz) = gbdrift0_file
         end do
      end do
      cvdrift0 = gbdrift0

      close (geofile_unit)

    end subroutine overwrite_selected_geometric_coefficients

  end subroutine init_geometry

  subroutine allocate_arrays (nalpha, nzgrid)

    implicit none

    integer, intent (in) :: nalpha, nzgrid

    if (.not.allocated(bmag)) allocate (bmag(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(bmag_psi0)) allocate (bmag_psi0(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds2)) allocate (gds2(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds21)) allocate (gds21(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds22)) allocate (gds22(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds23)) allocate (gds23(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds24)) allocate (gds24(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds25)) allocate (gds25(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gds26)) allocate (gds26(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dgds2dr))  allocate (dgds2dr(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dgds21dr)) allocate (dgds21dr(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dgds22dr)) allocate (dgds22dr(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gbdrift)) allocate (gbdrift(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(gbdrift0)) allocate (gbdrift0(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(cvdrift)) allocate (cvdrift(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(cvdrift0)) allocate (cvdrift0(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dgbdriftdrho)) allocate (dgbdriftdrho(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dcvdriftdrho)) allocate (dcvdriftdrho(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dgbdrift0drho)) allocate (dgbdrift0drho(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dcvdrift0drho)) allocate (dcvdrift0drho(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dbdzed)) allocate (dbdzed(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(theta_vmec)) allocate (theta_vmec(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(jacob)) allocate (jacob(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(djacdrho)) allocate (djacdrho(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(grho)) allocate (grho(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(dl_over_b)) allocate (dl_over_b(nalpha,-nzgrid:nzgrid))
    if (.not.allocated(d_dl_over_b_drho)) allocate (d_dl_over_b_drho(nalpha,-nzgrid:nzgrid))

    if (.not.allocated(gradpar)) allocate (gradpar(-nzgrid:nzgrid))
    if (.not.allocated(zed_eqarc)) allocate (zed_eqarc(-nzgrid:nzgrid))
    if (.not.allocated(btor)) allocate (btor(-nzgrid:nzgrid))
    if (.not.allocated(rmajor)) allocate (rmajor(-nzgrid:nzgrid))
    if (.not.allocated(dBdrho)) allocate (dBdrho(-nzgrid:nzgrid))
    if (.not.allocated(d2Bdrdth)) allocate (d2Bdrdth(-nzgrid:nzgrid))
    if (.not.allocated(dgradpardrho)) allocate (dgradpardrho(-nzgrid:nzgrid))

    if (.not.allocated(alpha)) allocate (alpha(nalpha)) ; alpha = 0.
    if (.not.allocated(zeta)) allocate (zeta(nalpha,-nzgrid:nzgrid)) ; zeta = 0.

  end subroutine allocate_arrays

  subroutine read_parameters

    use text_options
    use file_utils, only: error_unit, input_unit_exist
    use file_utils, only: runtype_option_Switch, runtype_multibox
    use mp, only: job
    use physics_flags, only: radial_variation

    implicit none

    integer :: in_file, ierr
    logical :: exist

    character (20) :: geo_option
    type (text_option), dimension (5), parameter :: geoopts = (/ &
         text_option('default', geo_option_local), &
         text_option('miller', geo_option_local), &
         text_option('local', geo_option_local), &
         text_option('input.profiles', geo_option_inputprof), &
         text_option('vmec', geo_option_vmec) /)

    namelist /geo_knobs/ geo_option, geo_file, overwrite_bmag, overwrite_gradpar, &
         overwrite_gds2, overwrite_gds21, overwrite_gds22, overwrite_gds23, overwrite_gds24, &
         overwrite_gbdrift, overwrite_cvdrift, overwrite_gbdrift0, q_as_x

    geo_option = 'local'
    overwrite_bmag = .false.
    overwrite_gradpar = .false.
    overwrite_gds2 = .false.
    overwrite_gds21 = .false.
    overwrite_gds22 = .false.
    overwrite_gds23 = .false.
    overwrite_gds24 = .false.
    overwrite_gbdrift = .false.
    overwrite_cvdrift = .false.
    overwrite_gbdrift0 = .false.
    q_as_x = radial_variation !true by default in radial variation runs
    geo_file = 'input.geometry'

    in_file = input_unit_exist("geo_knobs", exist)
    if (exist) read (unit=in_file, nml=geo_knobs)

    ierr = error_unit()
    call get_option_value &
         (geo_option, geoopts, geo_option_switch, &
         ierr, "geo_option in geo_knobs")

    if(radial_variation.and.runtype_option_switch.eq.runtype_multibox.and.job.ne.1) then
      geo_option_switch = geo_option_multibox
    endif

    overwrite_geometry = overwrite_bmag .or. overwrite_gradpar &
         .or. overwrite_gds2 .or. overwrite_gds21 .or. overwrite_gds22 &
         .or. overwrite_gds23 .or. overwrite_gds24 &
         .or. overwrite_cvdrift .or. overwrite_gbdrift .or. overwrite_gbdrift0

  end subroutine read_parameters

  subroutine broadcast_arrays

    use mp, only: broadcast

    implicit none

    call broadcast (qinp)
    call broadcast (shat)
    call broadcast (drhodpsi)
    call broadcast (drhodpsi_psi0)
    call broadcast (exb_nonlin_fac)
    call broadcast (exb_nonlin_fac_p)
    call broadcast (dIdrho)
    call broadcast (grho)
    call broadcast (bmag)
    call broadcast (bmag_psi0)
    call broadcast (btor)
    call broadcast (rmajor)
    call broadcast (gradpar)
    call broadcast (gds2)
    call broadcast (gds21)
    call broadcast (gds22)
    call broadcast (gds23)
    call broadcast (gds24)
    call broadcast (gds25)
    call broadcast (gds26)
    call broadcast (dgds2dr)
    call broadcast (dgds21dr)
    call broadcast (dgds22dr)
    call broadcast (gbdrift0)
    call broadcast (gbdrift)
    call broadcast (cvdrift0)
    call broadcast (cvdrift)
    call broadcast (dgbdrift0drho)
    call broadcast (dgbdriftdrho)
    call broadcast (dcvdrift0drho)
    call broadcast (dcvdriftdrho)
    call broadcast (dBdrho)
    call broadcast (d2Bdrdth)
    call broadcast (dgradpardrho)
    call broadcast (djacdrho)

    call broadcast (geo_surf%rmaj)
    call broadcast (geo_surf%rgeo)
    call broadcast (geo_surf%kappa)
    call broadcast (geo_surf%kapprim)
    call broadcast (geo_surf%tri)
    call broadcast (geo_surf%triprim)
    call broadcast (geo_surf%rhoc)
    call broadcast (geo_surf%rhoc_psi0)
    call broadcast (geo_surf%dr)
    call broadcast (geo_surf%shift)
    call broadcast (geo_surf%qinp)
    call broadcast (geo_surf%qinp_psi0)
    call broadcast (geo_surf%shat)
    call broadcast (geo_surf%shat_psi0)
    call broadcast (geo_surf%betaprim)
    call broadcast (geo_surf%betadbprim)
    call broadcast (geo_surf%d2qdr2)
    call broadcast (geo_surf%d2psidr2)
    call broadcast (geo_surf%dpsitordrho)
    call broadcast (geo_surf%rhotor)
    call broadcast (geo_surf%psitor_lcfs)
    call broadcast (geo_surf%drhotordrho)

    call broadcast (zed_scalefac)
    call broadcast (q_as_x)
    call broadcast (alpha)
    call broadcast (zeta)
    call broadcast (dxdXcoord)
    call broadcast (dydalpha)
    call broadcast (twist_and_shift_geo_fac)

    call broadcast (aref)
    call broadcast (bref)

  end subroutine broadcast_arrays

  subroutine communicate_geo_multibox(l_edge, r_edge)

    use millerlocal, only: communicate_parameters_multibox
    use mp, only: proc0

    implicit none

    real, intent (in) :: l_edge, r_edge

    if(proc0) then
      call communicate_parameters_multibox(geo_surf,gfac*l_edge, gfac*r_edge)
    endif

  end subroutine communicate_geo_multibox

  ! given function f(z:-pi->pi), calculate z derivative
  ! second order accurate, with equal grid spacing assumed
  ! assumes periodic in z -- may need to change this in future
  subroutine get_dzed (nz, dz, f, df)

    implicit none

    integer, intent (in) :: nz
    real, dimension (-nz:), intent (in) :: dz, f
    real, dimension (-nz:), intent (out) :: df

    df(-nz+1:nz-1) = (f(-nz+2:)-f(:nz-2))/(dz(:nz-2)+dz(-nz+1:nz-1))

    ! FLAG -- THIS MAY NEED TO BE CHANGED
    ! use periodicity at boundary
    df(-nz) = (f(-nz+1)-f(nz-1))/(dz(-nz)+dz(nz-1))
    df(nz) = df(-nz)

  end subroutine get_dzed

  subroutine get_gradpar_eqarc (gp, z, dz, gp_eqarc)

    use constants, only: pi
    use zgrid, only: nzgrid

    implicit none

    real, dimension (-nzgrid:), intent (in) :: gp, z, dz
    real, intent (out) :: gp_eqarc

    ! first get int dz b . grad z
    call integrate_zed (dz, 1./gp, gp_eqarc)
    ! then take (zmax-zmin)/int (dz b . gradz)
    ! to get b . grad z'
    gp_eqarc = (z(nzgrid)-z(-nzgrid))/gp_eqarc

  end subroutine get_gradpar_eqarc

  subroutine get_zed_eqarc (gp, dz, z, gp_eqarc, z_eqarc)

    use zgrid, only: nzgrid

    implicit none

    real, dimension (-nzgrid:), intent (in) :: gp, dz, z
    real, intent (in) :: gp_eqarc
    real, dimension (-nzgrid:), intent (out) :: z_eqarc

    integer :: iz

    z_eqarc(-nzgrid) = z(-nzgrid)
    do iz = -nzgrid+1, nzgrid
       call integrate_zed (dz(:iz), 1./gp(:iz), z_eqarc(iz))
    end do
    z_eqarc(-nzgrid+1:) = z(-nzgrid) + z_eqarc(-nzgrid+1:)*gp_eqarc

  end subroutine get_zed_eqarc

  ! trapezoidal rule to integrate in zed
  subroutine integrate_zed (dz, f, intf)

    use zgrid, only: nzgrid

    implicit none

    real, dimension (-nzgrid:), intent (in) :: dz
    real, dimension (-nzgrid:), intent (in) :: f
    real, intent (out) :: intf

    integer :: iz, iz_max

    iz_max = -nzgrid + size(dz) - 1

    intf = 0.
    do iz = -nzgrid+1, iz_max
       intf = intf + dz(iz)*(f(iz-1)+f(iz))
    end do
    intf = 0.5*intf

  end subroutine integrate_zed

  subroutine get_x_to_rho (llim, x_in, rho_out)

    use physics_parameters, only: rhostar

    implicit none

    integer, intent (in) :: llim

    real, dimension (:), intent (in) :: x_in
    real, dimension (:), intent (out) :: rho_out

    integer :: ix, ulim

    real :: a, b, c

    ulim = size(x_in)+llim-1

    if(q_as_x) then
      a = 0.5*geo_surf%d2qdr2/dqdrho
      b = 1.0
      c = -rhostar/(dqdrho*dxdXcoord)
    else
      a = 0.5*geo_surf%d2psidr2*drhodpsi
      b = 1.0
      c = -rhostar*drhodpsi/dxdXcoord
    endif

    do ix=llim, ulim
      if(abs(4.0*a*c*x_in(ix)) .lt. 1.e-6) then
        rho_out(ix) = -(c*x_in(ix))/b - a*(c*x_in(ix))**2/b**3
      else
        rho_out(ix) = (-b + sqrt(b**2 - 4.*a*c*x_in(ix)))/(2.*a)
      endif
    enddo

  end subroutine get_x_to_rho

  subroutine write_geometric_coefficients (nalpha)

    use file_utils, only: open_output_file, close_output_file
    use zgrid, only: nzgrid, zed

    implicit none

    integer, intent (in) :: nalpha
    integer :: geometry_unit
    integer :: ia, iz

    call open_output_file (geometry_unit,'.geometry')

    write (geometry_unit,'(a1,11a14)') '#', 'rhoc', 'qinp', 'shat', 'rhotor', 'aref', 'bref', 'dxdXcoord', 'dydalpha', &
                                      'exb_nonlin', 'exb_nonlin_p'
    write (geometry_unit,'(a1,11e14.4)') '#', geo_surf%rhoc, geo_surf%qinp, geo_surf%shat, geo_surf%rhotor, aref, bref, &
                                              dxdXcoord, dydalpha, exb_nonlin_fac, exb_nonlin_fac_p*exb_nonlin_fac
    write (geometry_unit,*)

    write (geometry_unit,'(15a12)') '# alpha', 'zed', 'zeta', 'bmag', 'gradpar', 'gds2', 'gds21', 'gds22', &
                                    'gds23', 'gds24', 'gbdrift', 'cvdrift', 'gbdrift0', 'bmag_psi0', 'btor'
    do ia = 1, nalpha
       do iz = -nzgrid, nzgrid
          write (geometry_unit,'(15e12.4)') alpha(ia), zed(iz), zeta(ia,iz), bmag(ia,iz), gradpar(iz), &
               gds2(ia,iz), gds21(ia,iz), gds22(ia,iz), gds23(ia,iz), &
               gds24(ia,iz), gbdrift(ia,iz), cvdrift(ia,iz), gbdrift0(ia,iz), &
               bmag_psi0(ia,iz), btor(iz)
       end do
       write (geometry_unit,*)
    end do

    call close_output_file (geometry_unit)

  end subroutine write_geometric_coefficients

  subroutine finish_init_geometry

    use mp, only: proc0
    use millerlocal, only: finish_local_geo

    implicit none

    if (proc0) then
       select case (geo_option_switch)
       case (geo_option_local)
         call finish_local_geo
       case (geo_option_multibox)
         call finish_local_geo
       case (geo_option_inputprof)
         call finish_local_geo
       case (geo_option_vmec)
       end select
    end if

  end subroutine finish_init_geometry

  subroutine finish_geometry

    implicit none

    if (allocated(zed_eqarc)) deallocate (zed_eqarc)
    if (allocated(grho)) deallocate (grho)
    if (allocated(bmag)) deallocate (bmag)
    if (allocated(bmag_psi0)) deallocate (bmag_psi0)
    if (allocated(btor)) deallocate (btor)
    if (allocated(rmajor)) deallocate (rmajor)
    if (allocated(dbdzed)) deallocate (dbdzed)
    if (allocated(jacob)) deallocate (jacob)
    if (allocated(djacdrho)) deallocate (djacdrho)
    if (allocated(gradpar)) deallocate (gradpar)
    if (allocated(dl_over_b)) deallocate (dl_over_b)
    if (allocated(d_dl_over_b_drho)) deallocate (d_dl_over_b_drho)
    if (allocated(gds2)) deallocate (gds2)
    if (allocated(gds21)) deallocate (gds21)
    if (allocated(gds22)) deallocate (gds22)
    if (allocated(gds23)) deallocate (gds23)
    if (allocated(gds24)) deallocate (gds24)
    if (allocated(gds25)) deallocate (gds25)
    if (allocated(gds26)) deallocate (gds26)
    if (allocated(dgds2dr))  deallocate (dgds2dr)
    if (allocated(dgds21dr)) deallocate (dgds21dr)
    if (allocated(dgds22dr)) deallocate (dgds22dr)
    if (allocated(gbdrift)) deallocate (gbdrift)
    if (allocated(gbdrift0)) deallocate (gbdrift0)
    if (allocated(cvdrift)) deallocate (cvdrift)
    if (allocated(cvdrift0)) deallocate (cvdrift0)
    if (allocated(dgbdriftdrho)) deallocate (dgbdriftdrho)
    if (allocated(dcvdriftdrho)) deallocate (dcvdriftdrho)
    if (allocated(dgbdrift0drho)) deallocate (dgbdrift0drho)
    if (allocated(dcvdrift0drho)) deallocate (dcvdrift0drho)
    if (allocated(dBdrho)) deallocate (dBdrho)
    if (allocated(d2Bdrdth)) deallocate (d2Bdrdth)
    if (allocated(dgradpardrho)) deallocate (dgradpardrho)
    if (allocated(theta_vmec)) deallocate (theta_vmec)

    if (allocated(alpha)) deallocate (alpha)
    if (allocated(zeta)) deallocate (zeta)

    geoinit = .false.

  end subroutine finish_geometry

end module stella_geometry


c=======================================================================
c
c     subroutine  CHECKSDLS
c
c     Computing the constraints matrices for SDLSCI
c
c-----------------------------------------------------------------------
      SUBROUTINE checksdls ( nmat, cmat, mcte, acte, bcte, mcti, acti,
     &                       lcti, ucti, dwork, info)
c-----------------------------------------------------------------------
c
c     INPUT
c            nmat   : dimension of the matrix                    integer
c            cmat   : symmetric matrix(nmat*nmat) to optimize     double
c            mcte   : number of equal constraints                integer
c            acte   : mcte Ai symmetric matrices(nmat*nmat) of
c                     equal constraints, vector(mcte*nmat*nmat)   double
c            bcte   : value of equal constraints, vector(mcte)    double
c            mcti   : number of inequal constraints              integer
c            acti   : mcti Bi symmetric matrices(nmat*nmat) of
c                     inequal constraints  vector(mcti*nmat*nmat) double
c            lcti   : lower bounds of inequal constraints
c                     vector(mcti)                                double
c            ucti   : upper bounds of inequal constraints
c                     vector(mcti)                                double
c
c     WORKSPACE:
c            dwork  : double workspace needs nmat*nmat            double
c
c     OUTPUT
c            info   : diagnostic argument                         integer
c
c     CALL
c       checksdlseq
c       checksdlsineq
c
c-----------------------------------------------------------------------
c
      IMPLICIT NONE
c
c     i/o arguments
      INTEGER nmat, mcte, mcti, info
      DOUBLE PRECISION cmat(nmat,*), acte(*), bcte(*), acti(*), lcti(*)
      DOUBLE PRECISION ucti(*)
c
c     Workspace
      double precision dwork(*)
c
c-----------------------------------------------------------------------
c
      CALL checksdlseq ( nmat, cmat, mcte, acte, bcte,
     &                   dwork, info )
      IF (info .LT. 0) RETURN
c
      CALL checksdlsineq ( nmat, cmat, mcti, acti, lcti, ucti,
     &                     dwork, info)
c
      RETURN
      END


!--------------------------------------------------------------------------------
! Copyright (c) 2016 Peter Grünberg Institut, Forschungszentrum Jülich, Germany
! This file is part of FLEUR and available as free software under the conditions
! of the MIT license as expressed in the LICENSE file in more detail.
!--------------------------------------------------------------------------------

MODULE m_prpqfftmap
  use m_juDFT
CONTAINS
  SUBROUTINE prp_qfft_map(stars,sym,input, igq2_fft,igq_fft)
    !*********************************************************************
    !     This subroutine prepares the pointer which identifies a
    !     threedimensional g-vector in the positive domain of the
    !     charge density fft-box in order to map a 3-d g-vector
    !     onto stars in case of the backtransform for fft of the
    !     charge density. correspondes  to igfft(*,2)
    !     it is independent of spin and k-point.
    !     pointer is built up when ever the chargedensity is calculated
    !     in order to save memory
    !
    !        s. bluegel, JRCAT, Feb. 97
    !*********************************************************************
    !
    USE m_types
    USE m_constants
    IMPLICIT NONE
    TYPE(t_input),INTENT(IN)   :: input
    TYPE(t_sym),INTENT(IN)     :: sym
    TYPE(t_stars),INTENT(IN)   :: stars
    !
    !
    INTEGER igq2_fft(0:stars%kq1_fft*stars%kq2_fft-1),igq_fft(0:stars%kq1_fft*stars%kq2_fft*stars%kq3_fft-1)
    !
    !---> local variables
    !
    LOGICAL new
    INTEGER istr,iop,iopm1,il,im,in,kid2x,kidx,iv1d,ifftq1,ifftq2
    INTEGER norm,kr(3,sym%nop),nop_local

    !------->          ABBREVIATIONS
    !
    !     kq1d  : dimension of the charge density FFT box in the pos. domain
    !     kq2d  : defined in dimens.f program (subroutine apws).1,2,3 indicate
    !     kq3d  ; a_1, a_2, a_3 directions.
    !     kq(i) : i=1,2,3 actual length of the fft-box for which FFT is done.
    !     nstr  : number of members (arms) of reciprocal lattice (g) vector
    !             of each star.
    !     nq3_fft: number of stars in the  charge density  FFT-box
    !     kmxq_fft: number of g-vectors forming the nq3_fft stars in the
    !               charge density sphere
    !
    !-----> prepare pointer which identifies a threedimensional g-vector
    !       in the positive domain of the charge density fft-box.
    !       correspondes  to igfft(*,2)
    !
    kidx    = 0
    kid2x   = 0
    ifftq1  = stars%kq1_fft
    ifftq2  = stars%kq1_fft*stars%kq2_fft
    !
    DO istr = 1 , stars%ng3_fft
       !
       nop_local=sym%nop
       IF (stars%kv3(3,istr).EQ.0) nop_local=sym%nop2
       !
       DO iop = 1,nop_local
          kr(1,iop) = stars%kv3(1,istr)*sym%mrot(1,1,iop) &
               + stars%kv3(2,istr)*sym%mrot(2,1,iop)+ stars%kv3(3,istr)*sym%mrot(3,1,iop)
          kr(2,iop) = stars%kv3(1,istr)*sym%mrot(1,2,iop) &
               + stars%kv3(2,istr)*sym%mrot(2,2,iop)+ stars%kv3(3,istr)*sym%mrot(3,2,iop)
          kr(3,iop) = stars%kv3(1,istr)*sym%mrot(1,3,iop) &
               + stars%kv3(2,istr)*sym%mrot(2,3,iop) + stars%kv3(3,istr)*sym%mrot(3,3,iop)
       ENDDO
       !
       DO iop = 1 , nop_local
          new=.true.
          DO iopm1 = 1 , iop - 1
             norm=(kr(1,iop)-kr(1,iopm1))**2 +&
                  (kr(2,iop)-kr(2,iopm1))**2 +(kr(3,iop)-kr(3,iopm1))**2
             IF (norm.EQ.0) new=.false.
          ENDDO

          IF (new) THEN
             il=kr(1,iop)
             im=kr(2,iop)
             in=kr(3,iop)
             if(il.lt.0) il=il+stars%kq1_fft
             if(im.lt.0) im=im+stars%kq2_fft
             if(in.lt.0) in=in+stars%kq3_fft
             iv1d = in*ifftq2 + im*ifftq1 + il
             igq_fft(kidx)=iv1d
             kidx=kidx+1
             IF (input%film.AND.(stars%kv3(3,istr).EQ.0)) THEN
                iv1d = im*ifftq1 + il
                igq2_fft(kid2x)=iv1d
                kid2x=kid2x+1
             ENDIF
          ENDIF
       ENDDO
       !
    ENDDO
    !
    IF (kidx .NE. stars%kmxq_fft) THEN
       WRITE (oUnit,'('' something wrong with stars%kmxq_fft or nq3_fft'')')
       WRITE (oUnit,'('' stars%kmxq_fft, acutal kidx '',2i5)') stars%kmxq_fft, kidx
       CALL juDFT_error("something wrong with stars or nq3_fft", calledby ="prp_qfft_map")
    ENDIF

  END SUBROUTINE prp_qfft_map
END MODULE m_prpqfftmap


c*************************************************************************
c                          ANAL_ENERGY_HJS.F
c*************************************************************************
c Returns the energy of the total system (massive bodies) in
c   the hierarchical Jacobi case
c
c      Input:
c            nbod          ==>  number of massive bodies (int scalar)
c            mass          ==>  mass of bodies (real array)
c            umat          ==>  Conversion matrix Jacobi => Barycentric
c                                (2D real array)
c            xj,yj,zj      ==>  current position in Jacobi coord
c                               (real arrays)
c            vxj,vyj,vzj   ==>  current velocity in Jacobi coord
c                               (real arrays)
c            eoff          ==> An energy offset that is added to the energy
c                                      (real*8 scalar)
c      Output:
c            energy        ==> the energy
c
c Remarks: Adapted from anal_energy_write_hjs
c Authors:  Herve Beust
c Date:    Sep 4, 2006

      subroutine anal_energy_hjs(nbod,umat,mass,xj,yj,zj,
     &                                       vxj,vyj,vzj,eoff,energy)

      include '../swift.inc'

c...  Inputs:
      integer nbod,iu
      real*8 mass(nbod),t,umat(NPLMAX,NPLMAX),eoff
      real*8 xj(nbod),yj(nbod),zj(nbod)
      real*8 vxj(nbod),vyj(nbod),vzj(nbod)

c...  Internals
      integer i1st
      real*8 energy,ke,pot
      real*8 xx,yy,zz,rr2
      real*8 xb(NPLMAX),yb(NPLMAX),zb(NPLMAX)
      real*8 vxb(NPLMAX),vyb(NPLMAX),vzb(NPLMAX)
      integer i,j

c----
c...  Executable code

c Compute initial ke,pot,energy

      call coord_g2b(nbod,umat,mass,xj,yj,zj,vxj,vyj,vzj,
     &      xb,yb,zb,vxb,vyb,vzb)

      ke = 0.5d0*mass(nbod)*(vxb(nbod)**2+vyb(nbod)**2+vzb(nbod)**2)
      pot = 0.d0

      do i = 1,nbod-1

         ke = ke + 0.5d0*mass(i)*(vxb(i)**2+vyb(i)**2+vzb(i)**2)
         do j = i+1,nbod

	    xx = xb(i) - xb(j)
	    yy = yb(i) - yb(j)
	    zz = zb(i) - zb(j)
	    rr2 = xx**2 + yy**2 + zz**2
	    pot = pot - mass(i)*mass(j)/(sqrt(rr2))

         enddo
      enddo

      energy = ke + pot + eoff

      return
      end     !anal_energy_hjs
c--------------------------------------------------------------------------


module damageNewLib
use funcAux
use damageMemoryManagementLib
use HyperModelsLib
use fibresModelsLib
use damageModelsNewLib

implicit none

contains

subroutine modifyStressAndTangent(Ppk,D, damagePar,Param,NDimE,damageMethod)
	real*8, intent(inout) :: Ppk(:,:),D(:,:,:,:)
	real*8, intent(in) :: damagePar(:) , Param(*)
	integer , intent(in) :: NdimE, damageMethod
	real*8 :: PpktenPpk(NdimE,NdimE,NdimE,NdimE) , damage, derDamage

	call decideDamageModel(damage, derDamage , damagePar,Param,damageMethod)

	call VotimesW(PpktenPpk,Ppk,Ppk)

	D = (1.0 - damage)*D - derDamage*PpktenPpk
	Ppk = (1.0 - damage)*Ppk

end subroutine

subroutine modifyStressAndTangent_fib(Sfib,DSfib, damagePar,Param,NDimE,damageMethod)
	real*8, intent(inout) :: Sfib(:),DSfib(:,:)
	real*8, intent(in) :: damagePar(:) , Param(*)
	integer , intent(in) :: NdimE, damageMethod
	real*8 :: SfibtenSfib(NdimE,NdimE) , damage, derDamage

	call decideDamageModel(damage, derDamage , damagePar,Param,damageMethod)

	call VotimesW(SfibtenSfib,Sfib,Sfib)

	DSfib = (1.0 - damage)*DSfib - derDamage*SfibtenSfib
	Sfib = (1.0 - damage)*Sfib

end subroutine

subroutine modifyStressAndTangent_fib_1D(sfib,dsfib, sfib_e, damagePar,Param, damageMethod)
	real*8, intent(inout) :: sfib,dsfib
	real*8, intent(in) :: damagePar(:) , Param(*), sfib_e
	integer , intent(in) :: damageMethod
	real*8 :: damage, derDamage

	call decideDamageModel(damage, derDamage , damagePar,Param,damageMethod)

	dsfib = (1.0 - damage)*dsfib - derDamage*sfib*sfib_e
	sfib = (1.0 - damage)*sfib

end subroutine

subroutine stateVariableUpdate_fib(qfib,damagePar,Param,matPar, constlaw)
	real*8, intent(in) :: qfib(:), damagePar(:), matPar(:)
	real*8, intent(inout) :: Param(*)
	integer , intent(in) :: constLaw
	type(damageState) :: sd, sdNew
	real*8 :: H, Hder, r, r0, rtheta, theta, ftrial,energy0, rAux

	theta = damagePar(1)
	r0 = damagePar(3)

	call strainEnergy_fib(energy0,qfib,matPar,constLaw)

	call sd%loadDamage(Param, 1, .false.)

	call fixInitialVariables(sd,sdNew,r0,Param)

	ftrial = dsqrt(2.0d0*energy0) - sd%rd

	if(ftrial <= 0.0d0) then
		sdNew%rd = sd%rd
		sdNew%rdDot = 0.0d0
		sdNew%qd = sd%qd
		sdNew%damage = sd%damage
!~ 		sd%derDamage = 0.0d0
		sdNew%derDamage = 0.0d0
	else
		sdNew%rd = sd%rd + ftrial
		sdNew%rdDot = ftrial

		rtheta = theta*sdNew%rd + (1.0d0 - theta)*sd%rd

		call getH(H,Hder,rtheta,damagePar)

		sdNew%qd = sd%qd + H*(sdNew%rd - sd%rd)

		if(sdNew%qd < 0.0d0) sdNew%qd = 1.0d-15

		sdNew%damage = 1.0d0 - sdNew%qd/sdNew%rd !! integrating

		sdNew%derDamage = (sdNew%qd - sdNew%rd*(H + theta*Hder*(sdNew%rd - sd%rd) ))/(sdNew%rd**3.0d0)

!~ 		sdNew%derDamage = (sdNew%qd - sdNew%rd*H)/(sdNew%rd**3.0d0)	! wrong
	end if

	call sdNew%putDamage(Param, 1)

end subroutine

subroutine stateVariableUpdate(F,damagePar,Param,matPar, constlaw)
	use fibresMod , only: Ipos_r0

	real*8, intent(in) :: F(:,:), damagePar(:), matPar(:)
	real*8, intent(inout) :: Param(*)
	integer , intent(in) :: constLaw
	type(damageState) :: sd, sdNew
	real*8 :: H, Hder, r, r0, rtheta, theta, ftrial,energy0, rAux

	theta = damagePar(1)
	r0 = damagePar(3)

	call strainEnergy(energy0,F,matPar,constLaw)

	call sd%loadDamage(Param, 1, .false.)

	call fixInitialVariables(sd,sdNew,r0,Param)

	ftrial = dsqrt(2.0d0*energy0) - sd%rd

	if(ftrial <= 0.0d0) then
		sdNew%rd = sd%rd
		sdNew%rdDot = 0.0d0
		sdNew%qd = sd%qd
		sdNew%damage = sd%damage
		sdNew%derDamage = 0.0d0
	else
		sdNew%rd = sd%rd + ftrial
		sdNew%rdDot = ftrial

		rtheta = theta*sdNew%rd + (1.0d0 - theta)*sd%rd

		call getH(H,Hder,rtheta,damagePar)

		sdNew%qd = sd%qd + H*(sdNew%rd - sd%rd)

		if(sdNew%qd < 0.0d0) sdNew%qd = 1.0d-15

		sdNew%damage = 1.0d0 - sdNew%qd/sdNew%rd !! integrating

		sdNew%derDamage = (sdNew%qd - sdNew%rd*(H + theta*Hder*(sdNew%rd - sd%rd) ))/(sdNew%rd**3.0d0)

!~ 		sdNew%derDamage = (sdNew%qd - sdNew%rd*H)/(sdNew%rd**3.0d0)	! wrong
	end if

	call sdNew%putDamage(Param, 1)

end subroutine

real*8 function getImplexDamage(rd,rddot,qd,damagePar) result(damageTilde)
	real*8, intent(in) :: rd,rddot,qd,damagePar(:)
	real*8 :: rdtilde, qdtilde, H, Hder

	rdtilde = rd + rdDot
	call getH(H,Hder,rd,damagePar)
	qdTilde = qd + H*(rdtilde - rd)
	damageTilde = 1.0d0 - qdTilde/rdTilde

end function

subroutine fixInitialVariables(sd,sdNew,r0,param)
	type(damageState) , intent(inout) :: sd, sdNew
	real*8 , intent(out) :: Param(*)
	real*8 , intent(in) :: r0

	if(r0>sd%rd) then !!! because r should begin in r0
		sdNew%rd = r0
		sd%rd = r0
		sdNew%qd = r0
		sd%qd = r0
		sdNew%derDamage = 0.0d0
		sd%derDamage = 0.0d0
		sdNew%damage = 0.0d0
		sd%damage = 0.0d0
		sdNew%rdDot = 0.0e0
		sd%rdDot = 0.0e0
		call sd%putDamageIntoOld(Param, 1)
		call sdNew%putDamage(Param, 1)
	end if

end subroutine

subroutine decideDamageModel(damage, derDamage , damagePar,Param,damageMethod)
	real*8, intent(out) :: damage, derDamage
	real*8, intent(in) :: damagePar(:) , Param(*)
	integer , intent(in) :: damageMethod

	type(damageState) :: sdNew , sd

	call sdNew%loadDamage(Param, 1, .true.)
	call sd%loadDamage(Param, 1, .false.)

	select case(damageMethod)

		case(1) !! implicit
			derDamage = sdNew%derDamage
			damage = sdNew%damage

		case(2) !! explicit
			derDamage = 0.0d0
			damage = sd%damage

		case(3) !! implex
			derDamage = 0.0d0
			damage = getImplexDamage(sd%rd,sd%rddot,sd%qd,damagePar)

		case(4) !! wrong case, put the one converging
			derDamage = sd%derDamage
			damage = sdNew%damage

		case default  !! no damage at all
			derDamage = 0.0d0
			damage = 0.0d0

	end select

end subroutine

end module


      SUBROUTINE FIXPDS(N,Y,IFLAG,ARCTOL,EPS,TRACE,A,NDIMA,NFE,
     $     ARCLEN,YP,YPOLD,QR,LENQR,PIVOT,PP,WORK,WT,PHI,P,
     $     PAR,IPAR)
C
C SUBROUTINE  FIXPDS  FINDS A FIXED POINT OR ZERO OF THE
C N-DIMENSIONAL VECTOR FUNCTION F(X), OR TRACKS A ZERO CURVE
C OF A GENERAL HOMOTOPY MAP RHO(A,X,LAMBDA).  FOR THE FIXED
C POINT PROBLEM F(X) IS ASSUMED TO BE A C2 MAP OF SOME BALL
C INTO ITSELF.  THE EQUATION  X = F(X)  IS SOLVED BY
C FOLLOWING THE ZERO CURVE OF THE HOMOTOPY MAP
C
C  LAMBDA*(X - F(X)) + (1 - LAMBDA)*(X - A)  ,
C
C STARTING FROM LAMBDA = 0, X = A.  THE CURVE IS PARAMETERIZED
C BY ARC LENGTH S, AND IS FOLLOWED BY SOLVING THE ORDINARY
C DIFFERENTIAL EQUATION  D(HOMOTOPY MAP)/DS = 0  FOR
C Y(S) = (X(S), LAMBDA(S)).
C
C FOR THE ZERO FINDING PROBLEM F(X) IS ASSUMED TO BE A C2 MAP
C SUCH THAT FOR SOME R > 0,  X*F(X) >= 0  WHENEVER NORM(X) = R.
C THE EQUATION  F(X) = 0  IS SOLVED BY FOLLOWING THE ZERO CURVE
C OF THE HOMOTOPY MAP
C
C   LAMBDA*F(X) + (1 - LAMBDA)*(X - A)
C
C EMANATING FROM LAMBDA = 0, X = A.
C
C  A  MUST BE AN INTERIOR POINT OF THE ABOVE MENTIONED BALLS.
C
C FOR THE CURVE TRACKING PROBLEM RHO(A,X,LAMBDA) IS ASSUMED TO
C BE A C2 MAP FROM E**M X E**N X [0,1) INTO E**N, WHICH FOR
C ALMOST ALL PARAMETER VECTORS A IN SOME NONEMPTY OPEN SUBSET
C OF E**M SATISFIES
C
C  RANK [D RHO(A,X,LAMBDA)/D LAMBDA , D RHO(A,X,LAMBDA)/DX] = N
C
C FOR ALL POINTS (X,LAMBDA) SUCH THAT RHO(A,X,LAMBDA)=0.  IT IS
C FURTHER ASSUMED THAT
C
C           RANK [ D RHO(A,X0,0)/DX ] = N  .
C
C WITH A FIXED, THE ZERO CURVE OF RHO(A,X,LAMBDA) EMANATING
C FROM  LAMBDA = 0, X = X0  IS TRACKED UNTIL  LAMBDA = 1  BY
C SOLVING THE ORDINARY DIFFERENTIAL EQUATION
C D RHO(A,X(S),LAMBDA(S))/DS = 0  FOR  Y(S) = (X(S), LAMBDA(S)),
C WHERE S IS ARC LENGTH ALONG THE ZERO CURVE.  ALSO THE HOMOTOPY
C MAP RHO(A,X,LAMBDA) IS ASSUMED TO BE CONSTRUCTED SUCH THAT
C
C              D LAMBDA(0)/DS > 0  .
C
C THIS CODE IS BASED ON THE ALGORITHM IN L. T. WATSON, A
C GLOBALLY CONVERGENT ALGORITHM FOR COMPUTING FIXED POINTS OF
C C2 MAPS, APPL. MATH. COMPUT., 5 (1979) 297-311.
C
C
C FOR THE FIXED POINT AND ZERO FINDING PROBLEMS, THE USER
C MUST SUPPLY A SUBROUTINE  F(X,V)  WHICH EVALUATES F(X) AT X
C AND RETURNS THE VECTOR F(X) IN V, AND A SUBROUTINE
C  FJACS(X,QR,LENQR,PIVOT)  WHICH EVALUATES THE (SYMMETRIC)
C JACOBIAN MATRIX OF F(X) AT X, AND RETURNS THE SYMMETRIC
C JACOBIAN MATRIX IN PACKED SKYLINE STORAGE FORMAT IN QR.  LENQR
C AND PIVOT DESCRIBE THE DATA STRUCTURE IN QR.  FOR THE CURVE
C TRACKING PROBLEM, THE USER MUST SUPPLY A SUBROUTINE
C  RHOA(V,LAMBDA,X,PAR,IPAR)  WHICH GIVEN (X,LAMBDA) RETURNS A
C PARAMETER VECTOR A IN V SUCH THAT RHO(A,X,LAMBDA)=0, AND A
C SUBROUTINE  RHOJS(A,LAMBDA,X,QR,LENQR,PIVOT,PP,PAR,IPAR)  WHICH
C RETURNS IN QR THE SYMMETRIC N X N JACOBIAN MATRIX [D RHO/DX]
C EVALUATED AT (A,X,LAMBDA) AND STORED IN PACKED SKYLINE FORMAT,
C AND RETURNS IN PP THE VECTOR -(D RHO/D LAMBDA) EVALUATED AT
C (A,X,LAMBDA).  LENQR AND PIVOT DESCRIBE THE DATA STRUCTURE IN QR.
C *** NOTE THE MINUS SIGN IN THE DEFINITION OF PP. ***
C
C
C FUNCTIONS AND SUBROUTINES DIRECTLY OR INDIRECTLY CALLED BY FIXPDS:
C  D1MACH , F (OR  RHOA ), FJACS (OR  RHOJS ), FODEDS , GMFADS ,
C  MFACDS , MULTDS , PCGDS , QIMUDS , ROOT , SINTRP , SOLVDS ,
C  STEPDS , AND THE BLAS FUNCTIONS  DAXPY , DCOPY , DDOT , DNRM2 ,
C  DSCAL , IDAMAX .  ONLY  D1MACH  CONTAINS MACHINE DEPENDENT
C  CONSTANTS.  NO OTHER MODIFICATIONS BY THE USER ARE REQUIRED.
C
C ***WARNING:  THIS SUBROUTINE IS GENERALLY MORE ROBUST THAN  FIXPNS
C AND  FIXPQS , BUT MAY BE SLOWER THAN THOSE SUBROUTINES BY A
C FACTOR OF TWO.
C
C
C ON INPUT:
C
C N  IS THE DIMENSION OF X, F(X), AND RHO(A,X,LAMBDA).
C
C Y  IS AN ARRRAY OF LENGTH  N + 1.  (Y(1),...,Y(N)) = A  IS THE
C    STARTING POINT FOR THE ZERO CURVE FOR THE FIXED POINT AND
C    ZERO FINDING PROBLEMS.  (Y(1),...,Y(N)) = X0  FOR THE CURVE
C    TRACKING PROBLEM.
C
C IFLAG  CAN BE -2, -1, 0, 2, OR 3.  IFLAG  SHOULD BE 0 ON THE
C    FIRST CALL TO  FIXPDS  FOR THE PROBLEM  X=F(X), -1 FOR THE
C    PROBLEM  F(X)=0, AND -2 FOR THE PROBLEM  RHO(A,X,LAMBDA)=0.
C    IN CERTAIN SITUATIONS  IFLAG  IS SET TO 2 OR 3 BY  FIXPDS,
C    AND  FIXPDS  CAN BE CALLED AGAIN WITHOUT CHANGING  IFLAG.
C
C ARCTOL  IS THE LOCAL ERROR ALLOWED THE ODE SOLVER WHEN
C    FOLLOWING THE ZERO CURVE.  IF  ARCTOL .LE. 0.0  ON INPUT
C    IT IS RESET TO  .5*DSQRT(EPS).  NORMALLY  ARCTOL  SHOULD
C    BE CONSIDERABLY LARGER THAN  EPS.
C
C EPS  IS THE LOCAL ERROR ALLOWED THE ODE SOLVER WHEN VERY
C    NEAR THE FIXED POINT(ZERO).  EPS  IS APPROXIMATELY THE
C    MIXED ABSOLUTE AND RELATIVE ERROR IN THE COMPUTED FIXED
C    POINT(ZERO).
C
C TRACE  IS AN INTEGER SPECIFYING THE LOGICAL I/O UNIT FOR
C    INTERMEDIATE OUTPUT.  IF  TRACE .GT. 0  THE POINTS COMPUTED ON
C    THE ZERO CURVE ARE WRITTEN TO I/O UNIT  TRACE .
C
C A(1:NDIMA) CONTAINS THE PARAMETER VECTOR  A .  FOR THE FIXED POINT
C    AND ZERO FINDING PROBLEMS, A  NEED NOT BE INITIALIZED BY THE
C    USER, AND IS ASSUMED TO HAVE LENGTH  N.  FOR THE CURVE
C    TRACKING PROBLEM, A  HAS LENGTH  NDIMA  AND MUST BE INITIALIZED
C    BY THE USER.
C
C NDIMA  IS THE DIMENSION OF  A , AND IS USED ONLY FOR THE CURVE
C    TRACKING PROBLEM.
C
C YP(1:N+1) IS A WORK ARRAY CONTAINING THE CURRENT TANGENT
C    VECTOR TO THE ZERO CURVE.
C
C YPOLD(1:N+1) IS A WORK ARRAY CONTAINING THE PREVIOUS TANGENT
C    VECTOR TO THE ZERO CURVE.
C
C QR(1:LENQR)  IS A WORK ARRAY CONTAINING THE (SYMMETRIC) JACOBIAN
C    MATRIX WITH RESPECT TO X, IN THE PACKED SKYLINE STORAGE FORMAT.
C
C LENQR  IS THE DIMENSION OF  QR .
C
C PIVOT(1:N+2), PP(1:N), AND WORK(1:6*(N+1)+LENQR) ARE ALL WORK
C    ARRAYS USED BY  FODEDS  TO CALCULATE THE TANGENT VECTOR YP.
C
C WT(1:N+1), PHI(1:N+1,1:16), AND P(1:N+1) ARE ALL WORK ARRAYS
C    USED BY THE ODE SUBROUTINE  STEPDS  .
C
C PAR(1:*) AND IPAR(1:*) ARE ARRAYS FOR (OPTIONAL) USER PARAMETERS,
C    WHICH ARE SIMPLY PASSED THROUGH TO THE USER WRITTEN SUBROUTINES
C    RHOA, RHOJS.
C
C Y, ARCTOL, EPS, ARCLEN, NFE, AND IFLAG SHOULD ALL BE
C VARIABLES IN THE CALLING PROGRAM.
C
C
C ON OUTPUT:
C
C N  AND  TRACE  ARE UNCHANGED.
C
C (Y(1),...,Y(N)) = X, Y(N+1) = LAMBDA, AND Y IS AN APPROXIMATE
C    ZERO OF THE HOMOTOPY MAP.  NORMALLY LAMBDA = 1 AND X IS A
C    FIXED POINT(ZERO) OF F(X).  IN ABNORMAL SITUATIONS LAMBDA
C    MAY ONLY BE NEAR 1 AND X IS NEAR A FIXED POINT(ZERO).
C
C IFLAG =
C  -2   CAUSES  FIXPDS  TO INITIALIZE EVERYTHING FOR THE PROBLEM
C       RHO(A,X,LAMBDA) = 0 (USE ON FIRST CALL).
C
C  -1   CAUSES  FIXPDS  TO INITIALIZE EVERYTHING FOR THE PROBLEM
C       F(X) = 0 (USE ON FIRST CALL).
C
C   0   CAUSES  FIXPDS  TO INITIALIZE EVERYTHING FOR THE PROBLEM
C       X = F(X) (USE ON FIRST CALL).
C
C   1   NORMAL RETURN.
C
C   2   SPECIFIED ERROR TOLERANCE CANNOT BE MET.  EPS HAS BEEN
C       INCREASED TO A SUITABLE VALUE.  TO CONTINUE, JUST CALL
C       FIXPDS  AGAIN WITHOUT CHANGING ANY PARAMETERS.
C
C   3   STEPDS  HAS BEEN CALLED 1000 TIMES.  TO CONTINUE, CALL
C       FIXPDS  AGAIN WITHOUT CHANGING ANY PARAMETERS.
C
C   4   JACOBIAN MATRIX DOES NOT HAVE FULL RANK AND/OR THE CONJUGATE
C       GRADIENT ITERATION FOR THE KERNEL OF THE JACOBIAN MATRIX
C       FAILED TO CONVERGE.  THE ALGORITHM HAS FAILED (THE ZERO
C       CURVE OF THE HOMOTOPY MAP CANNOT BE FOLLOWED ANY FURTHER).
C
C   5   EPS  (OR  ARCTOL ) IS TOO LARGE.  THE PROBLEM SHOULD BE
C       RESTARTED BY CALLING  FIXPDS  WITH A SMALLER  EPS  (OR
C       ARCTOL ) AND  IFLAG = 0 (-1, -2).
C
C   6   I - DF(X)  IS NEARLY SINGULAR AT THE FIXED POINT (DF(X) IS
C       NEARLY SINGULAR AT THE ZERO, OR  D RHO(A,X,LAMBDA)/DX  IS
C       NEARLY SINGULAR AT  LAMBDA = 1 ).  ANSWER MAY NOT BE
C       ACCURATE.
C
C   7   ILLEGAL INPUT PARAMETERS, A FATAL ERROR.
C
C ARCTOL = EPS AFTER A NORMAL RETURN (IFLAG = 1).
C
C EPS  IS UNCHANGED AFTER A NORMAL RETURN (IFLAG = 1).  IT IS
C    INCREASED TO AN APPROPRIATE VALUE ON THE RETURN IFLAG = 2.
C
C A  WILL (NORMALLY) HAVE BEEN MODIFIED.
C
C NFE  IS THE NUMBER OF FUNCTION EVALUATIONS (= NUMBER OF
C    JACOBIAN EVALUATIONS).
C
C ARCLEN  IS THE LENGTH OF THE PATH FOLLOWED.
C
C
      DOUBLE PRECISION AOLD,ARCLEN,ARCTOL,CURSW,CURTOL,EPS,
     1  EPSSTP,EPST,H,HOLD,S,S99,SA,SB,SOUT,SQNP1,XOLD,Y1SOUT
      INTEGER IFLAG,IFLAGC,ITER,IVC,J,JW,K,KGI,KOLD,
     1  KSTEPS,LCODE,LENQR,LIMIT,LIMITD,N,NDIMA,NFE,NFEC,NP1,TRACE
      LOGICAL START,CRASH,ST99
C
C *****  ARRAY DECLARATIONS.  *****
C
C ARRAYS NEEDED BY THE ODE SUBROUTINE  STEPDS .
      DOUBLE PRECISION Y(N+1),WT(N+1),PHI(N+1,16),P(N+1),YP(N+1),
     1     ALPHAS(12),W(12),G(13),GI(11)
      INTEGER IV(10)
C
C ARRAYS NEEDED BY  FIXPDS , FODEDS , AND  PCGDS .
      DOUBLE PRECISION YPOLD(N+1),A(N),QR(LENQR),PP(N),
     1     WORK(6*(N+1)+LENQR),PAR(1)
      INTEGER PIVOT(N+2),IPAR(1)
C
C *****  END OF DIMENSIONAL INFORMATION.  *****
C
      SAVE
      EXTERNAL FODEDS
C
C LIMITD  IS AN UPPER BOUND ON THE NUMBER OF STEPS.  IT MAY BE
C CHANGED BY CHANGING THE FOLLOWING PARAMETER STATEMENT:
      PARAMETER (LIMITD=1000)
C
C
C :  :  :  :  :  :  :  :  :  :  :  :  :  :  :  :  :  :  :  :  :
      IF (N .LE. 0  .OR.  EPS .LE. 0.0 ) IFLAG=7
      IF (IFLAG .GE. -2  .AND.  IFLAG .LE. 0) GO TO 10
      IF (IFLAG .EQ. 2) GO TO 35
      IF (IFLAG .EQ. 3) GO TO 30
C ONLY VALID INPUT FOR  IFLAG  IS -2, -1, 0, 2, 3.
      IFLAG=7
      RETURN
C
C *****  INITIALIZATION BLOCK.  *****
C
10    ARCLEN=0.0
      S=0.0
      IF (ARCTOL .LE. 0.0) ARCTOL=.5*SQRT(EPS)
      NFEC=0
      IFLAGC=IFLAG
      NP1=N+1
      SQNP1=SQRT(DBLE(NP1))
C
C SWITCH FROM THE TOLERANCE  ARCTOL  TO THE (FINER) TOLERANCE  EPS  IF
C THE CURVATURE OF ANY COMPONENT OF  Y  EXCEEDS  CURSW.
C
      CURSW=10.0
C
      ST99=.FALSE.
      START=.TRUE.
      CRASH=.FALSE.
      HOLD=1.0
      H=.1
      EPSSTP=ARCTOL
      KSTEPS=0
C SET INITIAL CONDITIONS FOR ORDINARY DIFFERENTIAL EQUATION.
      YPOLD(NP1)=1.0
      YP(NP1)=1.0
      Y(NP1)=0.0
      WORK(2*NP1)=0.0
      WORK(3*NP1)=0.0
      DO 20 J=1,N
        YPOLD(J)=0.0
        YP(J)=0.0
        WORK(NP1+J)=0.0
        WORK(2*NP1+J)=0.0
20    CONTINUE
C LOAD  A  FOR THE FIXED POINT AND ZERO FINDING PROBLEMS.
      IF (IFLAGC .GE. -1) THEN
        CALL DCOPY(N,Y,1,A,1)
      ENDIF
30    LIMIT=LIMITD
C
C *****  END OF INITIALIZATION BLOCK.  *****
C
C
C *****  MAIN LOOP.  *****
C
35    DO 150 ITER=1,LIMIT
      IF (Y(NP1) .LT. 0.0) THEN
40      ARCLEN=ARCLEN+S
        IFLAG=5
        RETURN
      ENDIF
50    IF (S .LE. 7.0*SQNP1) GO TO 80
C ARC LENGTH IS GETTING TOO LONG, THE PROBLEM WILL BE
C RESTARTED WITH A DIFFERENT  A  VECTOR.
      ARCLEN=ARCLEN+S
      S=0.0
60    START=.TRUE.
      CRASH=.FALSE.
C COMPUTE A NEW  A  VECTOR.
      IF (IFLAGC .EQ. -2) THEN
        DO 63 JW=1,NDIMA
          QR(JW)=A(JW)
63      CONTINUE
        CALL RHOA(A,Y(NP1),Y,PAR,IPAR)
        DO 65 JW=1,NDIMA
          AOLD=QR(JW)
C IF NEW AND OLD  A  DIFFER BY TOO MUCH, TRACKING SHOULD NOT CONTINUE.
          IF (ABS(A(JW)-AOLD) .GT. 1.0+ABS(AOLD)) THEN
            ARCLEN=ARCLEN+S
            IFLAG=5
            RETURN
          ENDIF
65      CONTINUE
      ELSE
        CALL F(Y,YP)
        DO 70 JW=1,N
          AOLD=A(JW)
          IF (IFLAGC .EQ. -1) THEN
            A(JW)=Y(NP1)*YP(JW)/(1.0 - Y(NP1)) + Y(JW)
          ELSE
            A(JW)=(Y(JW) - Y(NP1)*YP(JW))/(1.0 - Y(NP1))
          ENDIF
C IF NEW AND OLD  A  DIFFER BY TOO MUCH, TRACKING SHOULD NOT CONTINUE.
          IF (ABS(A(JW)-AOLD) .GT. 1.0+ABS(AOLD)) THEN
            ARCLEN=ARCLEN+S
            IFLAG=5
            RETURN
          ENDIF
70      CONTINUE
      ENDIF
      GO TO 100
80    IF (Y(NP1) .LE. .99  .OR. ST99) GO TO 100
C WHEN LAMBDA REACHES .99, THE PROBLEM WILL BE RESTARTED WITH
C A NEW  A  VECTOR.
90    ST99=.TRUE.
      EPSSTP=EPS
      ARCTOL=EPS
      GO TO 60
C
C SET DIFFERENT ERROR TOLERANCE FOR HIGH CURVATURE COMPONENTS OF THE
C TRAJECTORY Y(S).
100   CURTOL=CURSW*HOLD
      EPST=EPS/EPSSTP
      DO 110 JW=1,NP1
        IF (ABS(YP(JW)-YPOLD(JW)) .LE. CURTOL) THEN
          WT(JW)=(ABS(Y(JW))+1.0)
        ELSE
          WT(JW)=(ABS(Y(JW))+1.0)*EPST
        ENDIF
110   CONTINUE
C
C TAKE A STEP ALONG THE CURVE.
      CALL STEPDS(FODEDS,NP1,Y,S,H,EPSSTP,WT,START,HOLD,K,KOLD,CRASH,
     +     PHI,P,YP,ALPHAS,W,G,KSTEPS,XOLD,IVC,IV,KGI,GI,
     +     YPOLD,A,QR,LENQR,PIVOT,PP,WORK,NFEC,IFLAGC,PAR,IPAR)
C PRINT LATEST POINT ON CURVE IF REQUESTED.
      IF (TRACE .GT. 0) THEN
        WRITE (TRACE,117) ITER,NFEC,S,Y(NP1),(Y(JW),JW=1,N)
117     FORMAT(/' STEP',I5,3X,'NFE =',I5,3X,'ARC LENGTH =',F9.4,3X,
     $  'LAMBDA =',F7.4,5X,'X vector:'/1P,(1X,6E12.4))
      ENDIF
      NFE=NFEC
C CHECK IF THE STEP WAS SUCCESSFUL.
      IF (IFLAGC .EQ. 4) THEN
        ARCLEN=ARCLEN+S
        IFLAG=4
        RETURN
      ENDIF
120   IF (CRASH) THEN
C RETURN CODE FOR ERROR TOLERANCE TOO SMALL.
        IFLAG=2
C CHANGE ERROR TOLERANCES.
        EPS=EPSSTP
        IF (ARCTOL .LT. EPS) ARCTOL=EPS
C CHANGE LIMIT ON NUMBER OF ITERATIONS.
        LIMIT=LIMIT-ITER
        RETURN
      ENDIF
C
130   IF (Y(NP1) .GE. 1.0) THEN
        IF (ST99) GO TO 160
C
C IF LAMBDA .GE. 1.0 BUT THE PROBLEM HAS NOT BEEN RESTARTED
C WITH A NEW  A  VECTOR, BACK UP AND RESTART.
C
        S99=S-.5*HOLD
C GET AN APPROXIMATE ZERO Y(S) WITH  Y(NP1)=LAMBDA .LT. 1.0  .
135     CALL SINTRP(S,Y,S99,WT,YP,NP1,KOLD,PHI,IVC,IV,KGI,GI,
     $     ALPHAS,G,W,XOLD,P)
        IF (WT(NP1) .LT. 1.0) GO TO 140
        S99=.5*(S-HOLD+S99)
        GO TO 135
C
140     CALL DCOPY(NP1,WT,1,Y,1)
        CALL DCOPY(NP1,YP,1,YPOLD,1)
        S=S99
        GO TO 90
      ENDIF
C
150   CONTINUE
C
C *****  END OF MAIN LOOP.  *****
C
C LAMBDA HAS NOT REACHED 1 IN 1000 STEPS.
      IFLAG=3
      RETURN
C
C
C USE INVERSE INTERPOLATION TO GET THE ANSWER AT LAMBDA = 1.0 .
C
160   SA=S-HOLD
      SB=S
      LCODE=1
170   CALL ROOT(SOUT,Y1SOUT,SA,SB,EPS,EPS,LCODE)
C ROOT  FINDS S SUCH THAT Y(NP1)(S) = LAMBDA = 1 .
      IF (LCODE .GT. 0) GO TO 190
      CALL SINTRP(S,Y,SOUT,WT,YP,NP1,KOLD,PHI,IVC,IV,KGI,GI,
     $     ALPHAS,G,W,XOLD,P)
      Y1SOUT=WT(NP1)-1.0
      GO TO 170
190   IFLAG=1
C SET IFLAG = 6 IF  ROOT  COULD NOT GET  LAMBDA = 1.0  .
      IF (LCODE .GT. 2) IFLAG=6
      ARCLEN=ARCLEN+SA
C LAMBDA(SA) = 1.0 .
      CALL SINTRP(S,Y,SA,WT,YP,NP1,KOLD,PHI,IVC,IV,KGI,GI,
     $     ALPHAS,G,W,XOLD,P)
C
      CALL DCOPY(NP1,WT,1,Y,1)
C
      RETURN
      END


module spherical_mod

!Balaji: parallel spectral model using transpose method
!initialize spherical operates on global domain
!all other operators attempt to detect whether argument is global (0:M,0:N)
!                                                       or local (ms:me,ns:ne)

use       fms_mod, only: mpp_pe, mpp_root_pe, error_mesg, FATAL, &
                         write_version_number

use  spec_mpp_mod, only: get_spec_domain

!-------------------------------------------------------------------------
!   provides operations on spectral spherical harmonics fields that do not
!      require transforms
!
!   spectral fields are complex with horizontal dimensions
!        (0:num_fourier,0:num_spherical)
!       where the zonal wavenumber of (m,n) is M = m*fourier_inc
!             the "meridional" wavenumber is n
!             the "total", 2D, spherical wavenumber L = M+n
!
!-----------------------------------------------------------------------

implicit none
private

character(len=128), parameter :: version = '$Id: spherical.f90,v 13.0 2006/03/28 21:18:30 fms Exp $'
character(len=128), parameter :: tagname ='$Name: latest $'

interface compute_lon_deriv_cos
   module procedure compute_lon_deriv_cos_2d,  &
                    compute_lon_deriv_cos_3d
end interface

interface compute_lat_deriv_cos
   module procedure compute_lat_deriv_cos_2d,  &
                    compute_lat_deriv_cos_3d
end interface

interface compute_gradient_cos
   module procedure compute_gradient_cos_2d,  &
                    compute_gradient_cos_3d
end interface

interface compute_laplacian
   module procedure compute_laplacian_2d,  &
                    compute_laplacian_3d
end interface

interface compute_ucos_vcos
   module procedure compute_ucos_vcos_2d,  &
                    compute_ucos_vcos_3d
end interface

interface compute_vor_div
   module procedure compute_vor_div_2d,  &
                    compute_vor_div_3d
end interface

interface compute_vor
   module procedure compute_vor_2d,  &
                    compute_vor_3d
end interface

interface compute_div
   module procedure compute_div_2d,  &
                    compute_div_3d
end interface

interface compute_alpha_operator
   module procedure compute_alpha_operator_2d,  &
                    compute_alpha_operator_3d
end interface

interface triangular_truncation
   module procedure triangular_truncation_2d,  &
                    triangular_truncation_3d
end interface

interface rhomboidal_truncation
   module procedure rhomboidal_truncation_2d,  &
                    rhomboidal_truncation_3d
end interface

public :: spherical_init, spherical_end
public :: compute_lon_deriv_cos, compute_lat_deriv_cos
public :: compute_laplacian, compute_vor, compute_div
public :: get_spherical_wave, get_fourier_wave
public :: get_eigen_laplacian, compute_gradient_cos
public :: compute_ucos_vcos, compute_vor_div
public :: triangular_truncation, rhomboidal_truncation

real,    allocatable, dimension(:,:) :: eigen_laplacian, epsilon, r_epsilon
integer, allocatable, dimension(:,:) :: fourier_wave, spherical_wave

integer :: fourier_max, fourier_inc, num_fourier, num_spherical
integer :: ms, me, ns, ne
logical :: module_is_initialized = .false.

real, allocatable, dimension(:,:) :: coef_uvm
real, allocatable, dimension(:,:) :: coef_uvc
real, allocatable, dimension(:,:) :: coef_uvp
real, allocatable, dimension(:,:) :: coef_alpm
real, allocatable, dimension(:,:) :: coef_alpp
real, allocatable, dimension(:,:) :: coef_dym
real, allocatable, dimension(:,:) :: coef_dx
real, allocatable, dimension(:,:) :: coef_dyp
real, allocatable, dimension(:,:) :: triangle_mask

contains

!--------------------------------------------------------------------------
subroutine spherical_init(radius, num_fourier_in, fourier_inc_in, num_spherical_in)

real,    intent(in) :: radius
integer, intent(in) :: num_fourier_in, fourier_inc_in, num_spherical_in

integer :: m,n

if(module_is_initialized) return

call write_version_number(version, tagname)

call get_spec_domain(ms, me, ns, ne)

num_fourier = num_fourier_in
fourier_inc = fourier_inc_in
fourier_max = fourier_inc*num_fourier
num_spherical = num_spherical_in

allocate(eigen_laplacian(0:num_fourier,0:num_spherical))
allocate(epsilon        (0:num_fourier,0:num_spherical))
allocate(r_epsilon      (0:num_fourier,0:num_spherical))
allocate(fourier_wave   (0:num_fourier,0:num_spherical))
allocate(spherical_wave (0:num_fourier,0:num_spherical))
allocate(coef_uvm       (0:num_fourier,0:num_spherical))
allocate(coef_uvc       (0:num_fourier,0:num_spherical))
allocate(coef_uvp       (0:num_fourier,0:num_spherical))
allocate(coef_alpm      (0:num_fourier,0:num_spherical))
allocate(coef_alpp      (0:num_fourier,0:num_spherical))
allocate(coef_dym       (0:num_fourier,0:num_spherical))
allocate(coef_dx        (0:num_fourier,0:num_spherical))
allocate(coef_dyp       (0:num_fourier,0:num_spherical))
allocate(triangle_mask  (0:num_fourier,0:num_spherical))

module_is_initialized = .true.

triangle_mask = 1.0
do n=0,num_spherical
  do m=0,num_fourier
    fourier_wave(m,n)   = m*fourier_inc
    spherical_wave(m,n) = fourier_wave(m,n) + n
    if(spherical_wave(m,n).gt.num_spherical-1) triangle_mask(m,n) = 0.0
  end do
end do

epsilon= (spherical_wave**2 - fourier_wave**2)/(4.0*spherical_wave**2 - 1.0)
epsilon = sqrt(epsilon)

eigen_laplacian = spherical_wave*(spherical_wave + 1.0)/(radius*radius)

where (spherical_wave > 0)
  coef_uvm = -radius*epsilon/spherical_wave
  coef_uvc = -radius*fourier_wave/(spherical_wave*(spherical_wave + 1.0))
else where
  coef_uvm = 0.0
  coef_uvc = 0.0
end where

coef_uvp(:,0:num_spherical-1) =    &
    -radius*epsilon(:,1:num_spherical)/   &
     (spherical_wave(:,0:num_spherical-1) +1.0)

coef_alpm= (spherical_wave + 1.0)*epsilon/radius
coef_alpp(:,0:num_spherical-1) =    &
     spherical_wave(:,0:num_spherical-1)*epsilon(:,1:num_spherical)/radius

coef_dym = (spherical_wave - 1.0)*epsilon/radius
coef_dx = float(fourier_wave)/radius
coef_dyp(:,0:num_spherical-1) =                                           &
  (spherical_wave(:,0:num_spherical-1) + 2.0)*epsilon(:,1:num_spherical)/radius

return
end subroutine spherical_init

!---------------------------------------------------------------------------
subroutine get_spherical_wave(spherical_wave_out)
!---------------------------------------------------------------------------

integer, intent(out), dimension (:,:) :: spherical_wave_out

if( size(spherical_wave_out,1).EQ.num_fourier+1 .AND. size(spherical_wave_out,2).EQ.num_spherical+1 )then
    spherical_wave_out = spherical_wave
else if( size(spherical_wave_out,1).EQ.me-ms+1 .AND. size(spherical_wave_out,2).EQ.ne-ns+1 )then
    spherical_wave_out = spherical_wave(ms:me,ns:ne)
else
    call error_mesg( 'get_spherical_wave', 'invalid argument size', FATAL )
endif

return
end subroutine get_spherical_wave

!---------------------------------------------------------------------------
subroutine get_fourier_wave(fourier_wave_out)
!---------------------------------------------------------------------------

integer, intent(out), dimension (:,:) :: fourier_wave_out

if( size(fourier_wave_out,1).EQ.num_fourier+1 .AND. size(fourier_wave_out,2).EQ.num_spherical+1 )then
    fourier_wave_out = fourier_wave
else if( size(fourier_wave_out,1).EQ.me-ms+1 .AND. size(fourier_wave_out,2).EQ.ne-ns+1 )then
    fourier_wave_out = fourier_wave(ms:me,ns:ne)
else
    call error_mesg( 'get_fourier_wave', 'invalid argument size', FATAL )
endif

return
end subroutine get_fourier_wave

!---------------------------------------------------------------------------
subroutine get_eigen_laplacian(eigen_laplacian_out)
!---------------------------------------------------------------------------

real, intent(out), dimension (:,:) :: eigen_laplacian_out

if( size(eigen_laplacian_out,1).EQ.num_fourier+1 .AND. size(eigen_laplacian_out,2).EQ.num_spherical+1 )then
    eigen_laplacian_out = eigen_laplacian
else if( size(eigen_laplacian_out,1).EQ.me-ms+1 .AND. size(eigen_laplacian_out,2).EQ.ne-ns+1 )then
    eigen_laplacian_out = eigen_laplacian(ms:me,ns:ne)
else
    call error_mesg( 'get_eigen_laplacian', 'invalid argument size', FATAL )
endif

return
end subroutine get_eigen_laplacian

!---------------------------------------------------------------------------
function compute_lon_deriv_cos_3d(spherical) result(deriv_lon)
!---------------------------------------------------------------------------

complex, intent(in), dimension (:,:,:) :: spherical
complex, dimension (size(spherical,1), size(spherical,2), size(spherical,3)) :: deriv_lon

integer :: k

if(.not. module_is_initialized ) then
  call error_mesg('compute_lon_deriv','module spherical not initialized', FATAL)
end if

if( size(spherical,1).EQ.num_fourier+1 .AND. size(spherical,2).EQ.num_spherical+1 )then
    do k=1,size(spherical,3)
       deriv_lon(:,:,k) = coef_dx(:,:)*cmplx(-aimag(spherical(:,:,k)),real(spherical(:,:,k)))
    end do
else if( size(spherical,1).EQ.me-ms+1 .AND. size(spherical,2).EQ.ne-ns+1 )then
    do k=1,size(spherical,3)
       deriv_lon(:,:,k) = coef_dx(ms:me,ns:ne)*cmplx(-aimag(spherical(:,:,k)),real(spherical(:,:,k)))
    end do
else
    call error_mesg( 'compute_lon_deriv_cos_3d', 'invalid argument size', FATAL )
endif

return
end function compute_lon_deriv_cos_3d

!--------------------------------------------------------------------------
function compute_lat_deriv_cos_3d(spherical) result(deriv_lat)
!--------------------------------------------------------------------------

complex, intent(in), dimension (:,0:,:) :: spherical
complex, dimension (size(spherical,1), 0:size(spherical,2)-1, size(spherical,3)) :: deriv_lat

integer :: k

if(.not. module_is_initialized ) then
  call error_mesg('compute_lat_deriv','module spherical not initialized', FATAL)
end if

if( size(spherical,2).EQ.num_spherical+1 )then
!could be global domain, or only global in N
    if( size(spherical,1).EQ.num_fourier+1 )then
        deriv_lat(:,0,:) = cmplx(0.,0.)
        do k=1,size(spherical,3)
           deriv_lat(:,1:num_spherical,k) = &
                - spherical(:,0:num_spherical-1,k)*coef_dym(:,1:num_spherical)
           deriv_lat(:,0:num_spherical-1,k) =  deriv_lat(:,0:num_spherical-1,k)   &
                + spherical(:,1:num_spherical,k)*coef_dyp(:,0:num_spherical-1)
        end do
    else if( size(spherical,1).EQ.me-ms+1 )then
        deriv_lat(:,0,:) = cmplx(0.,0.)
        do k=1,size(spherical,3)
           deriv_lat(:,1:num_spherical,k) = &
                - spherical(:,0:num_spherical-1,k)*coef_dym(ms:me,1:num_spherical)
           deriv_lat(:,0:num_spherical-1,k) =  deriv_lat(:,0:num_spherical-1,k)   &
                + spherical(:,1:num_spherical,k)*coef_dyp(ms:me,0:num_spherical-1)
        end do
    endif
else if( size(spherical,1).EQ.me-ms+1 .AND. size(spherical,2).EQ.ne-ns+1 )then
!need to write stuff to acquire data at ns-1,ne+1
    call abort()
else
    call error_mesg( 'compute_lat_deriv_cos_3d', 'invalid argument size', FATAL )
endif

return
end function compute_lat_deriv_cos_3d

!-------------------------------------------------------------------
subroutine compute_gradient_cos_3d(spherical, deriv_lon, deriv_lat)
!-------------------------------------------------------------------

complex, intent(in), dimension (:,:,:) :: spherical
complex, intent(out), dimension (:,:,:) :: deriv_lat
complex, intent(out), dimension (:,:,:) :: deriv_lon

deriv_lon = compute_lon_deriv_cos(spherical)
deriv_lat = compute_lat_deriv_cos(spherical)

return
end subroutine compute_gradient_cos_3d

!----------------------------------------------------------------------
function compute_laplacian_3d(spherical, power) result(laplacian)
!----------------------------------------------------------------------

  complex, intent(in), dimension (:,:,:) :: spherical
  integer, optional :: power

  complex, dimension (size(spherical,1), size(spherical,2), size(spherical,3)) :: laplacian

  integer :: k
  real, dimension(size(spherical,1), size(spherical,2)) :: factor

  if(.not. module_is_initialized ) then
      call error_mesg('compute_laplacian','module spherical not initialized', FATAL)
  end if

  if( size(spherical,1).EQ.num_fourier+1 .AND. size(spherical,2).EQ.num_spherical+1 )then
      if(present(power)) then
          if(power >= 0) then
              factor = (-eigen_laplacian)**power
          else
              where (eigen_laplacian .ne. 0.0)
                  factor = (-eigen_laplacian)**power
              else where
                  factor = 0.0
              end where
          end if
      else
          factor = -eigen_laplacian
      end if
  else if( size(spherical,1).EQ.me-ms+1 .AND. size(spherical,2).EQ.ne-ns+1 )then
      if(present(power)) then
          if(power >= 0) then
              factor = (-eigen_laplacian(ms:me,ns:ne))**power
          else
              where (eigen_laplacian(ms:me,ns:ne) .ne. 0.0)
                  factor = (-eigen_laplacian(ms:me,ns:ne))**power
              else where
                  factor = 0.0
              end where
          end if
      else
          factor = -eigen_laplacian(ms:me,ns:ne)
      end if
  else
      call error_mesg( 'compute_laplacian', 'invalid argument size', FATAL )
  endif

  do k= 1,size(spherical,3)
     laplacian(:,:,k) = spherical(:,:,k)*factor
  end do

  return
end function compute_laplacian_3d

!----------------------------------------------------------------------
subroutine compute_ucos_vcos_3d(vorticity , divergence, u_cos, v_cos)
!----------------------------------------------------------------------

complex, intent(in), dimension (:,0:,:) :: vorticity
complex, intent(in), dimension (:,0:,:) :: divergence
complex, intent(out), dimension (:,0:,:) :: u_cos
complex, intent(out), dimension (:,0:,:) :: v_cos

integer :: k

if(.not. module_is_initialized ) then
  call error_mesg('compute_ucos_vcos','module spherical not initialized', FATAL)
end if

if( size(vorticity,2).EQ.num_spherical+1 )then
!could be global domain, or only global in N
    if( size(vorticity,1).EQ.num_fourier+1 )then
        do k=1,size(vorticity,3)
           u_cos(:,:,k) = coef_uvc(:,:)*                                     &
                cmplx(-aimag(divergence(:,:,k)),real(divergence(:,:,k)))
           v_cos(:,:,k) = coef_uvc(:,:)*                                     &
                cmplx(-aimag(vorticity(:,:,k)),real(vorticity(:,:,k)))

           u_cos(:,1:num_spherical,k) = u_cos(:,1:num_spherical,k) +         &
                coef_uvm(:,1:num_spherical)*vorticity(:,0:num_spherical-1,k)
           v_cos(:,1:num_spherical,k) = v_cos(:,1:num_spherical,k) -         &
                coef_uvm(:,1:num_spherical)*divergence(:,0:num_spherical-1,k)

           u_cos(:,0:num_spherical-1,k) = u_cos(:,0:num_spherical-1,k) -     &
                coef_uvp(:,0:num_spherical-1)*vorticity(:,1:num_spherical,k)
           v_cos(:,0:num_spherical-1,k) = v_cos(:,0:num_spherical-1,k) +     &
                coef_uvp(:,0:num_spherical-1)*divergence(:,1:num_spherical,k)
        end do
    else if( size(vorticity,1).EQ.me-ms+1 )then
        do k=1,size(vorticity,3)
           u_cos(:,:,k) = coef_uvc(ms:me,:)*                                     &
                cmplx(-aimag(divergence(:,:,k)),real(divergence(:,:,k)))
           v_cos(:,:,k) = coef_uvc(ms:me,:)*                                     &
                cmplx(-aimag(vorticity(:,:,k)),real(vorticity(:,:,k)))

           u_cos(:,1:num_spherical,k) = u_cos(:,1:num_spherical,k) +         &
                coef_uvm(ms:me,1:num_spherical)*vorticity(:,0:num_spherical-1,k)
           v_cos(:,1:num_spherical,k) = v_cos(:,1:num_spherical,k) -         &
                coef_uvm(ms:me,1:num_spherical)*divergence(:,0:num_spherical-1,k)

           u_cos(:,0:num_spherical-1,k) = u_cos(:,0:num_spherical-1,k) -     &
                coef_uvp(ms:me,0:num_spherical-1)*vorticity(:,1:num_spherical,k)
           v_cos(:,0:num_spherical-1,k) = v_cos(:,0:num_spherical-1,k) +     &
                coef_uvp(ms:me,0:num_spherical-1)*divergence(:,1:num_spherical,k)
        end do
    endif
else if( size(vorticity,1).EQ.me-ms+1 .AND. size(vorticity,2).EQ.ne-ns+1 )then
!need to write stuff to acquire data at ns-1,ne+1
    call abort()
else
    call error_mesg( 'compute_ucos_vcos', 'invalid argument size', FATAL )
endif

return
end subroutine compute_ucos_vcos_3d

!-------------------------------------------------------------------------
subroutine compute_vor_div_3d(u_cos, v_cos, vorticity, divergence)
!-------------------------------------------------------------------------

complex, intent(in), dimension (:,:,:) :: u_cos
complex, intent(in), dimension (:,:,:) :: v_cos
complex, intent(out), dimension (:,:,:) :: vorticity
complex, intent(out), dimension (:,:,:) :: divergence

vorticity  = compute_alpha_operator(v_cos, u_cos, -1)
divergence = compute_alpha_operator(u_cos, v_cos, +1)

return
end subroutine compute_vor_div_3d

!-------------------------------------------------------------------------
function compute_vor_3d(u_cos, v_cos) result(vorticity)
!-------------------------------------------------------------------------

complex, intent(in), dimension (:,:,:) :: u_cos
complex, intent(in), dimension (:,:,:) :: v_cos
complex, dimension (size(u_cos,1), size(u_cos,2), size(u_cos,3)) :: vorticity

vorticity = compute_alpha_operator(v_cos, u_cos, -1)

return
end function compute_vor_3d

!-------------------------------------------------------------------------
function compute_div_3d(u_cos, v_cos) result(divergence)
!-------------------------------------------------------------------------

complex, intent(in), dimension (:,:,:) :: u_cos
complex, intent(in), dimension (:,:,:) :: v_cos
complex, dimension (size(u_cos,1), size(u_cos,2), size(u_cos,3)) :: divergence

divergence = compute_alpha_operator(u_cos, v_cos, +1)

return
end function compute_div_3d

!--------------------------------------------------------------------------------
function compute_alpha_operator_3d(spherical_a, spherical_b, isign) result(alpha)
!--------------------------------------------------------------------------------

complex, intent(in), dimension (:,0:,:) :: spherical_a
complex, intent(in), dimension (:,0:,:) :: spherical_b
integer,intent(in) :: isign

complex, dimension (size(spherical_a,1), 0:size(spherical_a,2)-1, size(spherical_a,3)) :: alpha

integer :: k

if(.not. module_is_initialized ) then
  call error_mesg('compute_vor or div','module spherical not initialized', FATAL)
end if

alpha = cmplx(0.,0.)

if( size(spherical_a,2).EQ.num_spherical+1 )then
!could be global domain, or only global in N
    if( size(spherical_a,1).EQ.num_fourier+1 )then
        do k=1,size(spherical_a,3)
           alpha(:,:,k) = coef_dx(:,:)*    &
                cmplx(-aimag(spherical_a(:,:,k)),real(spherical_a(:,:,k)))
           alpha(:,1:num_spherical,k) = alpha(:,1:num_spherical,k) -  &
                isign*coef_alpm(:,1:num_spherical)  &
                *spherical_b(:,0:num_spherical-1,k)
           alpha(:,0:num_spherical-1,k) = alpha(:,0:num_spherical-1,k) +  &
                isign*coef_alpp(:,0:num_spherical-1)*spherical_b(:,1:num_spherical,k)
        end do
    else if( size(spherical_a,1).EQ.me-ms+1 )then
        do k=1,size(spherical_a,3)
           alpha(:,:,k) = coef_dx(ms:me,:)*    &
                cmplx(-aimag(spherical_a(:,:,k)),real(spherical_a(:,:,k)))
           alpha(:,1:num_spherical,k) = alpha(:,1:num_spherical,k) -  &
                isign*coef_alpm(ms:me,1:num_spherical)  &
                *spherical_b(:,0:num_spherical-1,k)
           alpha(:,0:num_spherical-1,k) = alpha(:,0:num_spherical-1,k) +  &
                isign*coef_alpp(ms:me,0:num_spherical-1)*spherical_b(:,1:num_spherical,k)
        end do
    endif
else if( size(spherical_a,1).EQ.me-ms+1 .AND. size(spherical_a,2).EQ.ne-ns+1 )then
!need to write stuff to acquire data at ns-1,ne+1
    call abort()
else
    call error_mesg( 'compute_alpha_operator_3d', 'invalid argument size', FATAL )
endif

return
end function compute_alpha_operator_3d

!-----------------------------------------------------------------------
subroutine triangular_truncation_3d(spherical, trunc)
!-----------------------------------------------------------------------

complex, intent(inout), dimension (:,:,:) :: spherical
integer, intent(in), optional :: trunc
integer :: k

if( size(spherical,1).EQ.num_fourier+1 .AND. size(spherical,2).EQ.num_spherical+1 )then
    if(present(trunc)) then
        do k=1, size(spherical,3)
           where (spherical_wave > trunc)
               spherical(:,:,k) = cmplx(0.,0.)
           end where
        end do
    else
        do k=1, size(spherical,3)
           spherical(:,:,k) = spherical(:,:,k)*triangle_mask
        end do
    end if
else if( size(spherical,1).EQ.me-ms+1 .AND. size(spherical,2).EQ.ne-ns+1 )then
    if(present(trunc)) then
        do k=1, size(spherical,3)
           where (spherical_wave(ms:me,ns:ne) > trunc)
               spherical(:,:,k) = cmplx(0.,0.)
           end where
        end do
    else
        do k=1, size(spherical,3)
           spherical(:,:,k) = spherical(:,:,k)*triangle_mask(ms:me,ns:ne)
        end do
    end if
else
    call error_mesg( 'triang_trunc', 'invalid argument size', FATAL )
endif

return
end subroutine triangular_truncation_3d

!-----------------------------------------------------------------------------
subroutine rhomboidal_truncation_3d(spherical, trunc_fourier, trunc_spherical)
!-----------------------------------------------------------------------------

complex, intent(inout), dimension (0:,0:,:) :: spherical
integer, intent(in), optional:: trunc_fourier, trunc_spherical

integer :: k,n,m

if( size(spherical,1).EQ.num_fourier+1 .AND. size(spherical,2).EQ.num_spherical+1 )then
    if(present(trunc_fourier).and.present(trunc_spherical)) then
        do k=1,size(spherical,3)
           do n=0,size(spherical,2)-1
              do m=0,size(spherical,1)-1
                 if(fourier_wave(m,n) .GT. trunc_fourier.OR.n .GT. trunc_spherical)  &
                      spherical(m,n,k) = cmplx(0.,0.)
              end do
           end do
        end do
    else
        spherical(:,num_spherical,:) = cmplx(0.,0.)
    end if
else if( size(spherical,1).EQ.me-ms+1 .AND. size(spherical,2).EQ.ne-ns+1 )then
!wavenumbers are numbered from (0,0) which is actually (ms,ns)
    if(present(trunc_fourier).and.present(trunc_spherical)) then
        do k=1,size(spherical,3)
           do n=0,size(spherical,2)-1
              do m=0,size(spherical,1)-1
                 if(fourier_wave(m+ms,n+ns).GT.trunc_fourier .OR. n+ns.GT.trunc_spherical)  &
                      spherical(m,n,k) = cmplx(0.,0.)
              end do
           end do
        end do
    else
!truncate on n=num_spherical, only on domain containing n=num_spherical
        if( ne.EQ.num_spherical )spherical(:,ne-ns,:) = cmplx(0.,0.)
    end if
else
    call error_mesg( 'rhomboid_trunc', 'invalid argument size', FATAL )
endif

return
end subroutine rhomboidal_truncation_3d

!rest are 2d versions of 3d routines

!---------------------------------------------------------------------------
function compute_lon_deriv_cos_2d(spherical) result(deriv_lon)
!---------------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: spherical
complex, dimension (size(spherical,1), size(spherical,2)) :: deriv_lon

complex, dimension (size(spherical,1), size(spherical,2), 1) :: spherical_3d
complex, dimension (size(spherical,1), size(spherical,2), 1) :: deriv_lon_3d

spherical_3d(:,:,1) = spherical(:,:)
deriv_lon_3d = compute_lon_deriv_cos_3d(spherical_3d)
deriv_lon(:,:) = deriv_lon_3d(:,:,1)

return
end function compute_lon_deriv_cos_2d

!--------------------------------------------------------------------------
function compute_lat_deriv_cos_2d(spherical) result(deriv_lat)
!--------------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: spherical
complex, dimension (size(spherical,1), size(spherical,2)) :: deriv_lat

complex, dimension (size(spherical,1), size(spherical,2), 1) :: spherical_3d
complex, dimension (size(spherical,1), size(spherical,2), 1) :: deriv_lat_3d

spherical_3d(:,:,1) = spherical(:,:)
deriv_lat_3d = compute_lat_deriv_cos_3d(spherical_3d)
deriv_lat(:,:) = deriv_lat_3d(:,:,1)

return
end function compute_lat_deriv_cos_2d

!-------------------------------------------------------------------
subroutine compute_gradient_cos_2d(spherical, deriv_lon, deriv_lat)
!-------------------------------------------------------------------

complex, intent(in),  dimension(:,:) :: spherical
complex, intent(out), dimension(size(spherical,1), size(spherical,2)) :: deriv_lat
complex, intent(out), dimension(size(spherical,1), size(spherical,2)) :: deriv_lon

complex, dimension(size(spherical,1), size(spherical,2), 1) :: spherical_3d
complex, dimension(size(spherical,1), size(spherical,2), 1) :: deriv_lat_3d
complex, dimension(size(spherical,1), size(spherical,2), 1) :: deriv_lon_3d

spherical_3d(:,:,1) = spherical(:,:)
call compute_gradient_cos_3d(spherical_3d, deriv_lon_3d, deriv_lat_3d)
deriv_lon(:,:) = deriv_lon_3d(:,:,1)
deriv_lat(:,:) = deriv_lat_3d(:,:,1)

return
end subroutine compute_gradient_cos_2d

!----------------------------------------------------------------------
function compute_laplacian_2d(spherical, power) result(laplacian)
!----------------------------------------------------------------------

complex, intent(in), dimension (:,:) :: spherical
integer, optional :: power

complex, dimension (size(spherical,1), size(spherical,2)) :: laplacian

complex, dimension(size(spherical,1), size(spherical,2), 1) :: spherical_3d
complex, dimension(size(spherical,1), size(spherical,2), 1) :: laplacian_3d

spherical_3d(:,:,1) = spherical(:,:)
laplacian_3d = compute_laplacian_3d(spherical_3d, power)
laplacian(:,:) = laplacian_3d(:,:,1)

return
end function compute_laplacian_2d

!----------------------------------------------------------------------
subroutine compute_ucos_vcos_2d(vorticity , divergence, u_cos, v_cos)
!----------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: vorticity
complex, intent(in),  dimension (size(vorticity,1), size(vorticity,2)) :: divergence
complex, intent(out), dimension (size(vorticity,1), size(vorticity,2)) :: u_cos
complex, intent(out), dimension (size(vorticity,1), size(vorticity,2)) :: v_cos

complex, dimension (size(vorticity,1), size(vorticity,2), 1) :: vorticity_3d
complex, dimension (size(vorticity,1), size(vorticity,2), 1) :: divergence_3d
complex, dimension (size(vorticity,1), size(vorticity,2), 1) :: u_cos_3d
complex, dimension (size(vorticity,1), size(vorticity,2), 1) :: v_cos_3d

vorticity_3d(:,:,1)  = vorticity(:,:)
divergence_3d(:,:,1) = divergence(:,:)
call compute_ucos_vcos_3d(vorticity_3d, divergence_3d, u_cos_3d, v_cos_3d)
u_cos(:,:) = u_cos_3d(:,:,1)
v_cos(:,:) = v_cos_3d(:,:,1)

return
end subroutine compute_ucos_vcos_2d

!-------------------------------------------------------------------------
subroutine compute_vor_div_2d(u_div_cos, v_div_cos, vorticity, divergence)
!-------------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: u_div_cos
complex, intent(in),  dimension (size(u_div_cos,1), size(u_div_cos,2)) :: v_div_cos
complex, intent(out), dimension (size(u_div_cos,1), size(u_div_cos,2)) :: vorticity
complex, intent(out), dimension (size(u_div_cos,1), size(u_div_cos,2)) :: divergence

complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: u_div_cos_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: v_div_cos_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: vorticity_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: divergence_3d

u_div_cos_3d(:,:,1) = u_div_cos(:,:)
v_div_cos_3d(:,:,1) = v_div_cos(:,:)
call compute_vor_div_3d(u_div_cos_3d, v_div_cos_3d, vorticity_3d, divergence_3d)
vorticity(:,:)  = vorticity_3d(:,:,1)
divergence(:,:) = divergence_3d(:,:,1)

return
end subroutine compute_vor_div_2d

!-------------------------------------------------------------------------
function compute_vor_2d(u_div_cos, v_div_cos) result(vorticity)
!-------------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: u_div_cos
complex, intent(in),  dimension (:,:) :: v_div_cos
complex, dimension (size(u_div_cos,1), size(u_div_cos,2)) :: vorticity

complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: u_div_cos_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: v_div_cos_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: vorticity_3d

u_div_cos_3d(:,:,1) = u_div_cos(:,:)
v_div_cos_3d(:,:,1) = v_div_cos(:,:)
vorticity_3d = compute_vor_3d(u_div_cos_3d, v_div_cos_3d)
vorticity(:,:) = vorticity_3d(:,:,1)

return
end function compute_vor_2d

!-------------------------------------------------------------------------
function compute_div_2d(u_div_cos, v_div_cos) result(divergence)
!-------------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: u_div_cos
complex, intent(in),  dimension (:,:) :: v_div_cos
complex, dimension (size(u_div_cos,1), size(u_div_cos,2)) :: divergence

complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: u_div_cos_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: v_div_cos_3d
complex, dimension (size(u_div_cos,1), size(u_div_cos,2), 1) :: divergence_3d

u_div_cos_3d(:,:,1) = u_div_cos(:,:)
v_div_cos_3d(:,:,1) = v_div_cos(:,:)
divergence_3d = compute_div_3d(u_div_cos_3d, v_div_cos_3d)
divergence(:,:) = divergence_3d(:,:,1)

return
end function compute_div_2d

!--------------------------------------------------------------------------------
function compute_alpha_operator_2d(spherical_a, spherical_b, isign) result(alpha)
!--------------------------------------------------------------------------------

complex, intent(in),  dimension (:,:) :: spherical_a
complex, intent(in),  dimension (:,:) :: spherical_b
integer, intent(in) :: isign

complex, dimension (size(spherical_a,1), size(spherical_a,2)) :: alpha

complex, dimension (size(spherical_a,1), size(spherical_a,2), 1) :: spherical_a_3d
complex, dimension (size(spherical_a,1), size(spherical_a,2), 1) :: spherical_b_3d
complex, dimension (size(spherical_a,1), size(spherical_a,2), 1) :: alpha_3d

spherical_a_3d(:,:,1) = spherical_a(:,:)
spherical_b_3d(:,:,1) = spherical_b(:,:)
alpha_3d = compute_alpha_operator_3d(spherical_a_3d, spherical_b_3d, isign)
alpha(:,:) = alpha_3d(:,:,1)

return
end function compute_alpha_operator_2d

!-----------------------------------------------------------------------
subroutine triangular_truncation_2d(spherical, trunc)
!-----------------------------------------------------------------------

complex, intent(inout), dimension (:,:) :: spherical
integer, intent(in), optional :: trunc
complex, dimension (size(spherical,1), size(spherical,2), 1) :: spherical_3d

spherical_3d(:,:,1) = spherical(:,:)
call triangular_truncation_3d(spherical_3d, trunc)
spherical(:,:) = spherical_3d(:,:,1)

return
end subroutine triangular_truncation_2d

!--------------------------------------------------------------------------------
subroutine rhomboidal_truncation_2d(spherical, trunc_fourier, trunc_spherical)
!--------------------------------------------------------------------------------

complex, intent(inout), dimension (:,:) :: spherical
integer, intent(in), optional :: trunc_fourier, trunc_spherical

complex, dimension (size(spherical,1), size(spherical,2), 1) :: spherical_3d

spherical_3d(:,:,1) = spherical(:,:)
call rhomboidal_truncation_3d(spherical_3d, trunc_fourier, trunc_spherical)
spherical(:,:) = spherical_3d(:,:,1)

return
end subroutine rhomboidal_truncation_2d

!-----------------------------------------------------------------------
subroutine spherical_end

if(.not.module_is_initialized) return

deallocate(eigen_laplacian, epsilon, r_epsilon, fourier_wave, spherical_wave)
deallocate(coef_uvm, coef_uvc, coef_uvp, coef_alpm, coef_alpp, coef_dym, coef_dx, coef_dyp, triangle_mask)
module_is_initialized = .false.

return
end subroutine spherical_end
!-----------------------------------------------------------------------

end module spherical_mod


c-----DMFIT, v 0.1
c
c     this test program tests the dmfit subroutine and function
c     it also shows an application to KK dark matter

      program test

      implicit none

      character*32 arg
c-----variables needed to run the routines: mass
      real*8 MX,z
c     annihilation channel, smoothing parameter, switch diff/integral (see below)
      integer CH,HSM,SWITCH,ii
c     remember that the following channels, corresponding to WIMP annihilation modes
c     are implemented:
c
c     1: e^+e^-
c     2: mu^+ mu^-
c     3: tau^+ tau^-
c     4: b \bar b
c     5: t \bar t
c     6: gluon gluon
c     7: W+ W-
c     8: Z Z
c     9: c \bar c

c     gamma ray energy, in GeV
      real*8 EE

c-----outputs:
      real*8 DNDE,DNDM

c-----the actual dmfit functions
      real*8 dmfit_de,dmfit_dm

c-----test case: a 100 GeV WIMP in b\bar b, same smoothing as DarkSUSY's default
      CALL getarg(1,arg)
      read(arg,*) MX
      CALL getarg(2,arg)
      read(arg,*) CH
c      MX=6.d0
c      CH=12
      HSM=0

c-----let's look at an energy of 10 GeV
      EE=0.2d0

      call dmfit_load ("new_gammamc_dif.dat")

c     now let's get the integrated flux

c$$$      DNDE=dmfit_de(MX,CH,EE)
c$$$      DNDM=dmfit_dm(MX,CH,EE)

      write(*,*) '# mass',MX
      write(*,*) '# channel',CH
c      write(*,*) 'energy',EE

c$$$      write(*,*)
c$$$      write(*,*) 'dN/dE:',DNDE
c$$$      write(*,*) 'd(dN/dE)/dM:',DNDM
      write(*,*) '# idx energy Log(E/MX) EdN/dE d(dN/dE)/dM'
      do ii=0,249
         z=(ii+0.5)/250.
         EE=MX*10.**(10*(z-1.))
         DNDE=dmfit_de(MX,CH,EE)
c          write(*,*) DNDE
c         DNDM=dmfit_dm(MX,CH,EE)
         write(*,*) ii, EE,log10(EE)-log10(MX),DNDE!, DNDM
      enddo

      stop
      end

c-----these are the routines needed to be included:

      include 'dmfit_load.f'
      include 'dmfit_func.f'
      include 'dmfit_comm.f'