module apex
!
! Calculate quantities needed to transform scalar fields between geographic 
! and geomagnetic coordinate systems. 
!
  use shr_kind_mod ,only : r8 => shr_kind_r8
  use geogrid      ,only: nlon,nlonp1,nlonp2,nlat,ylatg,ylong,dlong,&
                          jspole,jnpole
  use maggrid      ,only: nmlon,nmlonp1,nmlat,ylatm,ylonm,dlonm

  implicit none
  save

  integer ::             &
    ig(nmlonp1,nmlat),   & ! geog lon grid containing each geomag point
    jg(nmlonp1,nmlat)      ! geog lat grid containing each geomag point

  real(r8) ::              &
    wt(4,nmlonp1,nmlat)      ! interpolation weights for geo2mag
 
  real(r8),dimension(nmlonp1,nmlat) :: & ! geo lat,lon coords on mag grid
    gdlatdeg,   & ! geographic latitude of each magnetic grid point (deg)
    gdlondeg      ! geographic longitude of each magnetic grid point (deg)
!
! Variables on geographic grid needed by other modules must
! be allocated dynamically to be grid-independent (sub alloc_apex):
!
  integer,allocatable :: & ! (nlonp1,jspole:jnpole))
    im(:,:),             & ! geomag lon grid containing each geog point
    jm(:,:)                ! geomag lat grid containing each geog point

  real(r8),allocatable ::  & ! (nlonp1,jspole:jnpole)
    dim(:,:),              & ! fraction in lon for grid interp
    djm(:,:)                 ! fraction in lat for grid interp

  real(r8),allocatable ::  & ! (nlonp1,jspole:jnpole,3,2)
    dvec(:,:,:,:)            ! vectors from apxmall

  real(r8),allocatable ::  & ! (nlonp1,jspole:jnpole)
    dddarr(:,:),           & ! from apxmall
    be3arr(:,:)              ! from apxmall

  real(r8),allocatable ::  & ! (nlonp1,jspole:jnpole)
    alatm(:,:),            & ! geomagnetic latitude at each geographic grid point (radians)
    alonm(:,:),            & ! geomagnetic longitude at each geographic grid point (radians)
    xb(:,:),               & ! northward component of magnetic field
    yb(:,:),               & ! eastward component of magnetic field
    zb(:,:),               & ! downward component of magnetic field (gauss)
    bmod(:,:)                ! magnitude of magnetic field (gauss)
!
! rjac: scaled derivatives of geomagnetic coords wrt geographic coordinates.
! rjac(1,1) = cos(thetas)/cos(theta)*d(lamdas)/d(lamda)
! rjac(1,2) = cos(thetas)*d(lamdas)/d(theta)
! rjac(2,1) = 1./cos(theta)*d(thetas)/d(lamda)
! rjac(2,2) = d(thetas)/d(theta)
! where (lamda,theta) are geographic coordinates
!       (lamdas,thetas) are geomagnetic coordinates
!
    real(r8),allocatable :: &
      rjac(:,:,:,:) ! (nlon+1,jspole:jnpole,2,2)

  contains
!-----------------------------------------------------------------------
  subroutine get_apex(epoch,nlon,nlat,jspole,jnpole)
!
! This is called once per run from main.
!
  use params,only: re_dyn,h0,hs,iulog,dtr,rtd
!
! Args:
    integer,intent(in)  :: nlon,nlat,jspole,jnpole
    real(r8),intent(in) :: epoch
!
! Local:
    integer :: i,j,ier,jjm,jjg
    integer :: nlonp1,nlonp2
    integer,parameter :: nalt=2
    integer,save :: lwk
    real(r8),allocatable,save :: wk(:) ! wk(lwk)
    real(r8) :: real8

    real(r8) ::       &
      gplat(jnpole),  & ! (jnpole)
      gplon(nlon+2)     ! (nlonp2)

    real(r8) ::          &
      gpalt(nalt),       &
      dellat,dellon

    real(r8) :: rekm,h0km,alt,hr,ror03,glat,glon,&
      xlonmi,qdlon,qdlat,gdlon,gdlat,xlongi,frki,frkj
!
! Geographic grid in radians:
    real(r8) ::  &
      cslatg(jspole:jnpole), & ! cos(theta) (jspole:jnpole)
      snlatg(jspole:jnpole), & ! sin(theta) (jspole:jnpole)
      cslong(nlon+1),        & ! cos(lamda) (nlonp1)
      snlong(nlon+1)           ! sin(lamda) (nlonp1)
!
! Non-scalar arguments returned by apxmall:
    real(r8) ::           &
      b(3),bhat(3),       &
      d1(3),d2(3),d3(3),  &
      e1(3),e2(3),e3(3),  &
      f1(2),f2(2)
    real(r8) :: bmag,alon,xlatm,vmp,w,d,be3,si,sim,xlatqd,f

    real(r8),dimension(nlon+1,jspole:jnpole) :: &
      dmlat,   & ! dipole latitude corresponding to apex of field line
      rmag11,  & ! (a1.a1)/p*sin(i)*cos(thetas)
      rmagc,   & ! (a1.a2)/p*sin(i)
      rmag2,   & ! 1./bmod
      rmag22,  & ! (a2.a2)/p*sin(i)/cos(thetas)
      rjacd,   & ! determinant of rjac 
      p

    real(r8),dimension(nlon+1,nlat) :: &
      cslatm,  & ! cos(thetas)
      snlatm,  & ! sin(thetas)
      cslonm,  & ! cos(lamdas)
      snlonm     ! sin(lamdas)
!
! av  = the two magnetic vectors a1 and a2
! av1 = a1
! av2 = a2/cos(thetas)
!
    real(r8) :: av(nlon+1,jspole:jnpole,3,2)

    nlonp1 = nlon+1
    nlonp2 = nlon+2
!
! Allocate arrays that are needed by other modules:
    call alloc_apex
!
! Allocate local memory (is released at end of this routine):
    lwk= jnpole*nlonp2*nalt*6
    allocate(wk(lwk))

    gpalt(1) = 90._r8
    gpalt(2) = 170._r8

    real8 = real(jnpole-1)
    dellat = 180._r8/real8
    do j=1,jnpole
      gplat(j) = (j-1)*dellat - 90._r8
    enddo

    real8 = real(nlon)
    dellon = 360._r8/real8
    do i=1,nlonp2
      real8 = real(i)
      gplon(i) = (real8-1.5_r8)*dellon - 180._r8
    enddo

! Call Pete's apex lib (wrapper (apxntrp.f90), or legacy (apxntrp_legacy.f)):
!   call apxmka(iulog,epoch,gplat,gplon,gpalt,jnpole,nlonp2,nalt,wk,ier)
!   call apxmka_legacy(iulog,epoch,gplat,gplon,gpalt,jnpole,nlonp2,nalt,wk,lwk,ier)

! Call routine in apex_subs.F:
    call apxmka(iulog,epoch,gplat,gplon,gpalt,jnpole,nlonp2,nalt,wk,lwk,ier)

    if (ier /= 0) call shutdown('apxmka error')

    rekm = re_dyn*1.e-5_r8  ! earth radius (km)
    h0km = h0*1.e-5_r8
    alt  = hs*1.e-5_r8  ! modified apex reference altitude (km)
    hr   = alt
    ror03= ((rekm + alt)/(rekm + h0km))**3
!
! Loop over 2d geographic grid:
!
    do j=jspole,jnpole
      glat = ylatg(j)*rtd
      cslatg(j) = cos(ylatg(j)) 
      snlatg(j) = sin(ylatg(j))
      do i=1,nlonp1
        glon = ylong(i)*rtd
!
! Call Pete's apex lib (wrapper (apxntrp.f90), or legacy (apxntrp_legacy.f)):
!       call apxmall (                             & 
!         glat,glon,alt,hr, wk, lwk,               & !Inputs
!         b,bhat,bmag,si,                          & !Mag Fld
!         alon,                                    & !Apx Lon
!         xlatm,vmp,w,d,be3,sim,d1,d2,d3,e1,e2,e3, & !Mod Apx
!         xlatqd,f,f1,f2 , ier)                      !Qsi-Dpl
!
! Call routine in apex_subs.F:
        call apxmall (                             & 
          glat,glon,alt,hr, wk,                    & !Inputs
          b,bhat,bmag,si,                          & !Mag Fld
          alon,                                    & !Apx Lon
          xlatm,vmp,w,d,be3,sim,d1,d2,d3,e1,e2,e3, & !Mod Apx
          xlatqd,f,f1,f2 , ier)                      !Qsi-Dpl

        if (ier /= 0) call shutdown('apxmall error')

        alatm(i,j)    =  xlatm*dtr
        alonm(i,j)    =  alon *dtr
        xb   (i,j)    =  b(2)*1.e-5_r8 ! nT -> gauss
        yb   (i,j)    =  b(1)*1.e-5_r8 ! nT -> gauss
        zb   (i,j)    = -b(3)*1.e-5_r8 ! nT -> gauss
        bmod (i,j)    =  bmag*1.e-5_r8 ! nT -> gauss

        p    (i,j)     = 1.e+5_r8*abs(sim/si)/(bmag*d)
        av   (i,j,1,1) = d1(1)*P(I,J)*si
        av   (i,j,2,1) = d1(2)*P(I,J)*si
        av   (i,j,3,1) = d1(3)*P(I,J)*si
        av   (i,j,1,2) =-d2(1)*P(I,J)*si/sim
        av   (i,j,2,2) =-d2(2)*P(I,J)*si/sim
        av   (i,j,3,2) =-d2(3)*P(I,J)*si/sim
        rjac (i,j,1,1) =  f2(2)
        rjac (i,j,1,2) = -f2(1)
        rjac (i,j,2,1) = -f1(2)
        rjac (i,j,2,2) =  f1(1)
        rjacd(i,j)     = f
!
! Set up parameters for magnetic to geographic interpolation.
!
        xlonmi = (alonm(i,j) - ylonm(1))/dlonm
        real8 = real(nmlon)
        if (xlonmi < 0._r8) xlonmi = xlonmi + real8
        im(i,j) = xlonmi
        real8 = real(im(i,j))
        dim(i,j) = xlonmi - real8
        im(i,j) = im(i,j) + 1
        if (im(i,j) >= nmlonp1) im(i,j) = im(i,j) - nmlon
        alatm(i,j) = min(alatm(i,j),ylatm(nmlat))
        do jjm=2,nmlat
          if (alatm(i,j) > ylatm(jjm)) cycle
          jm(i,j) = jjm - 1
          djm(i,j) = (alatm(i,j) - ylatm(jm(i,j)))/  &
                     (ylatm(jjm) - ylatm(jm(i,j)))
          exit
        enddo
        if (j /= jspole .and. j /= jnpole) then
          dvec(i,j,1,1) = d1(1)
          dvec(i,j,2,1) = d1(2)
          dvec(i,j,3,1) = d1(3)
          dvec(i,j,1,2) = d2(1)
          dvec(i,j,2,2) = d2(2)
          dvec(i,j,3,2) = d2(3)
          dddarr(i,j)   = d
!
! Scale be3 from 130 km to a reference height of 90 km.
          be3arr(i,j)   = be3*ror03
!
          cslatm(i,j) = cos(alatm(i,j))
          snlatm(i,j) = sin(alatm(i,j))
          cslonm(i,j) = cos(alonm(i,j))
          snlonm(i,j) = sin(alonm(i,j))
        endif
      enddo ! i=1,nlonp1
    enddo ! j=jspole,jnpole
!
    do i = 1,nlonp1
      cslong(i) = cos(ylong(i))
      snlong(i) = sin(ylong(i))
    enddo
!
! Set up parameters for geographic to magnetic interpolation
    do i=1,nmlonp1
      qdlon = ylonm(i)*rtd
      do j=1,nmlat
        qdlat = ylatm(j)*rtd
!
! Convert from Quasi-Dipole to geographic coordinates.
! gdlat,gdlon are returned by apxq2g.
!
! Call Pete's apex lib (wrapper (apxntrp.f90), or legacy (apxntrp_legacy.f)):
!       call apxq2g(qdlat,qdlon,alt,wk,lwk,gdlat,gdlon,ier)
!       call apxq2g_legacy(qdlat,qdlon,alt,wk,gdlat,gdlon,ier)
!
! Call routine in apex_subs.F:
        call apxq2g(qdlat,qdlon,alt,wk,gdlat,gdlon,ier)

        if (ier /= 0) then
          write(iulog,"(i3,i3,i3)") 'apxparm: error from APXQ2G: ier=',ier, &
            ' i=',i,' j=',j
          call shutdown('apxq2g ier')
        endif
        gdlat = gdlat*dtr
        gdlon = gdlon*dtr
        xlongi = (gdlon - ylong(1))/dlong
        real8 = real(nlon)
        if (xlongi < 0._r8) xlongi = xlongi + real8
        ig(i,j) = xlongi
        real8 = real(ig(i,j))
        frki = xlongi - real8
        ig(i,j) = ig(i,j) + 1
        if (ig(i,j) >= nlonp1) ig(i,j) = ig(i,j) - nlon
        gdlat = min(gdlat,ylatg(jnpole))
        do jjg=1,jnpole
          if (gdlat > ylatg(jjg)) cycle
          jg(i,j) = jjg - 1
          frkj = (gdlat - ylatg(jg(i,j)))/(ylatg(jjg) - ylatg(jg(i,j)))
!
! 99/2/25b Add one to JG to account for the fact that AG in geo2mag has
!  a second (J) index starting at 1, while the second index of the
!  array in the calling arguments begins at 0.
!
          jg(i,j) = jg(i,j) + 1
          exit
        enddo
        wt(1,i,j) = (1._r8 - frki)*(1._r8 - frkj)
        wt(2,i,j) =          frki *(1._r8 - frkj)
        wt(3,i,j) =          frki *frkj
        wt(4,i,j) = (1._r8 - frki)*frkj
        gdlatdeg(i,j) = gdlat*rtd
        gdlondeg(i,j) = gdlon*rtd
      enddo ! j=1,nmlat
    enddo ! i=1,nmlonp1
!
! Release "local" memory:
    deallocate(wk)
  end subroutine get_apex
!-----------------------------------------------------------------------
  subroutine alloc_apex

    allocate(im (nlonp1,jspole:jnpole))
    allocate(jm (nlonp1,jspole:jnpole))
    allocate(dim(nlonp1,jspole:jnpole))
    allocate(djm(nlonp1,jspole:jnpole))

    allocate(xb   (nlonp1,jspole:jnpole))
    allocate(yb   (nlonp1,jspole:jnpole))
    allocate(zb   (nlonp1,jspole:jnpole))
    allocate(bmod (nlonp1,jspole:jnpole))
    allocate(alatm(nlonp1,jspole:jnpole))
    allocate(alonm(nlonp1,jspole:jnpole))

    allocate(dvec  (nlonp1,jspole:jnpole,3,2))
    allocate(dddarr(nlonp1,jspole:jnpole))
    allocate(be3arr(nlonp1,jspole:jnpole))

    allocate(rjac(nlon+1,jspole:jnpole,2,2))

  end subroutine alloc_apex
!-----------------------------------------------------------------------
end module apex