module heelis
  use shr_kind_mod  ,only: r8 => shr_kind_r8 ! 8-byte reals
  use cam_logfile   ,only: iulog
  use edyn_maggrid  ,only: nmlon,nmlonp1,nmlat,ylonm,ylatm
  use edyn_geogrid  ,only: nlat,nlon
  use edyn_params   ,only: dtr,rtd,pi
!
! phihm and pfrac are output of this module:
!
  use edyn_solve,only: & 
    nmlat0       ,& ! (nmlat+1)/2
    phihm        ,& ! output high-latitude potential (nmlonp1,nmlat)
    pfrac           ! output NH fraction of potential (nmlonp1,nmlat0)
  implicit none
  save
  private
!
! Auroral parameters (taken from aurora.F of timegcm):
! (dimension of 2 is for south,north hemispheres)
!
  real(r8) ::    &
    offa(2)     ,& ! offset of oval towards 0 MLT relative to magnetic pole (rad)
    offc(2)     ,& !
    dskofa(2)   ,& ! offset of oval in radians towards 18 MLT (f(By))
    dskofc(2)   ,& !
    phin(2)     ,& ! night convection entrance in MLT converted to radians (f(By))
    phid(2)     ,& ! dayside convection entrance in MLT converted to radians (f(By))
    psim(2)     ,& ! night convection entrance in MLT converted to radians (f(By))
    psie(2)     ,& !
    pcen(2)     ,& !
    phidp0(2)   ,& !
    phidm0(2)   ,& !
    phinp0(2)   ,& !
    phinm0(2)   ,& !
    theta0(2)   ,& ! convection reversal boundary in radians
    rr1(2)         !
 
  public heelis_model

  contains
!-----------------------------------------------------------------------
  subroutine heelis_model(iday,secs,ctpoten,sunlons)
!
! Driver for Heelis empirical model to calculate high-latitude potential.
!
! Args:
    integer,intent(in)  :: iday     ! day of year
    real(r8),intent(in) :: secs     ! ut in seconds
    real(r8),intent(in) :: ctpoten  ! cross-tail potential
    real(r8),intent(in) :: sunlons(nlat)  ! sun's location
!
! Local:
    real(r8),dimension(nmlonp1,nmlat) :: pfrac_plt
!
! Set auroral parameters:
!
    call heelis_init(ctpoten)
!
! 9/23/15 btf: The old call to sub colath is replaced by call
!   to calc_pfrac in dpie_coupling.F90. calc_pfrac is called
!   for either weimer05 or heelis.
!
! Calculate  the heelis potential phihm in geomagnetic coordinates:
! (potm calls sub flwv32)
!
    call potm(sunlons)

!   pfrac_plt = 0._r8
!   pfrac_plt(:,1:nmlat0) = pfrac(:,:)
!   call addfld('PHIHM',' ',' ',phihm,'mlon',1,nmlonp1,&
!     'mlat',1,nmlat,0)
!   call addfld('PFRAC',' ',' ',pfrac_plt,'mlon',1,nmlonp1,&
!     'mlat',1,nmlat,0)

  end subroutine heelis_model
!-----------------------------------------------------------------------
  subroutine heelis_init(ctpoten)
!
! Auroral parameters for Heelis (taken from aurora.F of timegcm):
!
! This is called at every timestep because ctpoten may change with time.
! Time-dependent ctpoten (kV) is read from TIMEGCM and WACCM input files, 
! unless it was provided as a constant by the user via namelist (namelist.F90).
!
! Args:
    real(r8),intent(in) :: ctpoten  ! cross-tail potential
!
! As in tiegcm1.9
    offa(:)   =  1._r8*dtr
    offc(:)   =  1._r8*dtr
    dskofa(:) =  0._r8
    dskofc(:) =  0._r8
    phin(:)   =  180._r8*dtr
    phid(:)   =  0._r8
    psim(:)   =  0.5_r8 * ctpoten * 1000._r8
    psie(:)   = -0.5_r8 * ctpoten * 1000._r8
    pcen(:)   =  0._r8
    phidp0(:) =  90._r8*dtr
    phidm0(:) =  90._r8*dtr
    phinp0(:) =  90._r8*dtr
    phinm0(:) =  90._r8*dtr
    rr1(:)    = -2.6_r8
    theta0(:) = (-3.80_r8+8.48_r8*(ctpoten**0.1875_r8))*dtr
!
! Current TIEGCM trunk (Sep, 2015) has asymmetry(?) and By effect as follows:
!     offc(isouth) = 1.1*dtr
!     offc(inorth) = 1.1*dtr
!     dskofc(isouth) = (-0.08 + 0.15*byloc)*dtr
!     dskofc(inorth) = (-0.08 - 0.15*byloc)*dtr
!     phid(isouth) = (9.39 + 0.21*byloc - 12.) * h2deg * dtr
!     phid(inorth) = (9.39 - 0.21*byloc - 12.) * h2deg * dtr
!     phin(isouth) = (23.50 + 0.15*byloc - 12.) * h2deg * dtr
!     phin(inorth) = (23.50 - 0.15*byloc - 12.) * h2deg * dtr
!     psim(:) =  0.44 * ctpoten * 1000.
!     psie(:) = -0.56 * ctpoten * 1000.
!     pcen(isouth) = (-0.168 + 0.027*byloc) * ctpoten * 1000.
!     pcen(inorth) = (-0.168 - 0.027*byloc) * ctpoten * 1000.
!     phidp0(:) = 90.*dtr
!     phidm0(:) = 90.*dtr
!     phinp0(:) = 90.*dtr
!     phinm0(:) = 90.*dtr
!     rr1(:) = -2.6

!   write(iulog,"('heelis_init: ctpoten=',e12.4,' theta0=',e12.4)") ctpoten,theta0(1)

  end subroutine heelis_init
!-----------------------------------------------------------------------
  subroutine potm(sunlons)
  use edyn_params ,only: pi_dyn       ! pi used in dynamo calculations
  use edyn_maggrid,only: ylonm,ylatm  ! magnetic grid lons, lats
!
! Calculate heelis potential in geomagnetic coordinates.
!
! Args:
    real(r8),intent(in) :: sunlons(nlat)
!
! Local:
    integer :: i,j
    real(r8),dimension(nmlon) :: dlat,dlon,ratio
    integer,dimension(nmlon) :: iflag
!
    ratio(:) = 1._r8
    do j=1,nmlat
      iflag(:) = 1 ! must be updated at each j
      dlat(:) = ylatm(j)
      dlon(:) = ylonm(1:nmlon)-sunlons(1)
!
! flwv32 returns single-level Heelis potential in geomag coords:
!
      if (abs(ylatm(j)) > pi_dyn/6._r8) then
        call flwv32(dlat,dlon,ratio,iflag,nmlon,phihm(:,j),j)
      else
        phihm(1:nmlon,j) = 0._r8
      endif
    enddo ! j=1,nmlat
!
! Periodic point:
    do j=1,nmlat
      phihm(nmlonp1,j) = phihm(1,j)
    enddo ! j=1,nmlat
  end subroutine potm
!-----------------------------------------------------------------------
  subroutine flwv32(dlat,dlon,ratio,iflag,nmlon,poten,mlat)
!     
! Calculate heelis potential at current magnetic latitude mlat.
!     
  use edyn_params      ,only: pi_dyn
! 
! Args:
    integer,intent(in) :: mlat,nmlon
    integer,intent(inout) :: iflag(nmlon)
    real(r8),dimension(nmlon),intent(in)  :: dlat,dlon,ratio
    real(r8),dimension(nmlon+1),intent(out) :: poten
!
! Local:
    integer :: i,n,ihem
    real(r8),parameter :: eps=1.e-6_r8
    real(r8) :: &
      pi2,pih,sinthr1,psi(8),phirc,sinth0, &
      ofda,cosofa(2),sinofa(2),aslona(2),  &
      ofdc,cosofc(2),sinofc(2),aslonc(2),  &
      phdpmx(2),phnpmx(2),phnmmx(2),phdmmx(2)
    real(r8),dimension(nmlon) :: sinlat,coslat,sinlon,coslon,alon, &
      colat,wk1,wk2,wk3,phifun,phifn2
    integer :: ifn(nmlon)
    real(r8) :: phi(nmlon,8)
!
    pi2 = 2.0_r8*pi_dyn
    pih = 0.5_r8*pi_dyn
    do n=1,2
      ofda = sqrt(offa(n)**2+dskofa(n)**2)
      cosofa(n) = cos(ofda)
      sinofa(n) = sin(ofda)
      aslona(n) = asin(dskofa(n)/ofda)
!
      ofdc = sqrt(offc(n)**2+dskofc(n)**2)
      cosofc(n) = cos(ofdc)
      sinofc(n) = sin(ofdc)
      aslonc(n) = asin(dskofc(n)/ofdc)
!
      if (phin(n) < phid(n)) phin(n) = phin(n)+pi2  ! modifies aurora phin
      phdpmx(n) = .5_r8*min(pi,(phin(n)-phid(n)))
      phnpmx(n) = .5_r8*min(pi,(phid(n)-phin(n)+pi2))
      phnmmx(n) = phdpmx(n)
      phdmmx(n) = phnpmx(n)
    enddo ! n=1,2
!
! Set ihem=1,2 for South,North hemisphere: 
!
    ihem = int(dlat(max0(1,nlon/2))*2._r8/3.1416_r8+2._r8)
    sinth0 = sin(theta0(ihem))
!
! Average amie results show r1=-2.6 for 11.3 degrees
!   (0.1972 rad) beyond theta0.
!
    sinthr1 = sin(theta0(ihem)+0.1972_r8)
    psi(1) = psie(ihem)
    psi(3) = psim(ihem)
    do n=2,4,2
      psi(n) = psi(n-1)
    enddo ! n=2,4,2
    do n=1,4
      psi(n+4) = psi(n)
    enddo ! n=1,4
!
! Transform to auroral circle coordinates:
!
    do i=1,nmlon
      sinlat(i) = sin(abs(dlat(i)))
      coslat(i) = cos(dlat(i))
      sinlon(i) = sin(dlon(i)+aslonc(ihem))
      coslon(i) = cos(dlon(i)+aslonc(ihem))
      colat(i) = cosofc(ihem)*sinlat(i)-sinofc(ihem)*coslat(i)* &
        coslon(i) 
      alon(i) = mod(atan2(sinlon(i)*coslat(i),sinlat(i)* &
        sinofc(ihem)+cosofc(ihem)*coslat(i)*coslon(i))- &
        aslonc(ihem)+3._r8*pi_dyn,pi2)-pi_dyn
      colat(i) = acos(colat(i))*sqrt(ratio(i))
!
! Boundaries for longitudinal function:
!
      wk1(i) = ((colat(i)-theta0(ihem))/theta0(ihem))**2
      phi(i,4)=phid(ihem)+eps-min(phidm0(ihem)+wk1(i)* &
        (pih-phidm0(ihem)),phdmmx(ihem))
      phi(i,5)=phid(ihem)-eps+min(phidp0(ihem)+wk1(I)* &
        (pih-phidp0(ihem)),phdpmx(ihem))
      phi(i,6)=phin(ihem)+eps-min(phinm0(ihem)+wk1(i)* &
        (pih-phinm0(ihem)),phnmmx(ihem))
      phi(i,7)=phin(ihem)-eps+min(phinp0(ihem)+wk1(i)* &
        (pih-phinp0(ihem)),phnpmx(ihem))
      phi(i,1)=phi(i,5)-pi2
      phi(i,2)=phi(i,6)-pi2
      phi(i,3)=phi(i,7)-pi2
      phi(i,8)=phi(i,4)+pi2
      phifun(i)=0._r8
      phifn2(i) = 0._r8
      if (colat(i)-theta0(ihem) >= 0._r8) then
        ifn(i) = 3
      else
        ifn(i) = 2
      endif
      if (iflag(i) == 1) iflag(i) = ifn(i)
!
! Add ring current rotation to potential (phirc)
!
      phirc = 0._r8
      wk2(i) = mod(alon(i)+phirc+2._r8*pi2+pi_dyn,pi2)-pi_dyn
      wk3(i) = mod(alon(i)+phirc+3._r8*pi2,pi2)-pi_dyn
    enddo ! i=1,nmlon
!
! Longitudinal variation:
!
    do n=1,7
      do i=1,nmlon
        phifun(i)=phifun(i)+.25_r8*(psi(n)+psi(n+1)+(psi(n)-       &
          psi(n+1))*cos(mod(pi_dyn*(wk2(i)-phi(i,n))/(phi(i,n+1)-  &
          phi(i,n)),pi2)))*(1._r8-sign(1._r8,(wk2(i)-phi(i,n))*    &
          (wk2(i)-phi(i,n+1))))
        phifn2(i)=phifn2(i)+.25_r8*(psi(n)+psi(n+1)+(psi(n)-       &
          psi(n+1))*cos(mod(pi_dyn*(wk3(i)-phi(i,n))/(phi(i,n+1)-  &
          phi(i,n)),pi2)))*(1._r8-sign(1._r8,(wk3(i)-phi(i,n))*    &
          (wk3(i)-phi(i,n+1))))
      enddo
    enddo
!
! Evaluate total potential:
!

!   write(iulog,*) ' MINVAL(phifun), MAXVAL(phifun) ', MINVAL(phifun), MAXVAL(phifun)
!   write(iulog,*) ' MINVAL(colat), MAXVAL(colat) ', MINVAL(colat), MAXVAL(colat)
!   write(iulog,*) ' sinth0,ihem,rr1(ihem),sinthr1 ', sinth0,ihem,rr1(ihem),sinthr1

!   write(iulog,"('  rr1(ihem)=',e12.4,' pcen(ihem)=',e12.4,' rr1(ihem)=',e12.4)") &
!     theta0(ihem),pcen(ihem),rr1(ihem)
!   write(iulog,"('  sinth0=',e12.4,' sinthr1=',e12.4)") sinth0,sinthr1
!   write(iulog,"('  iflag=',/,(20i3))") iflag
!   write(iulog,"('  dlat =',/,(6e12.4))") dlat
!   write(iulog,"('  dlon =',/,(6e12.4))") dlon
!   write(iulog,"('  colat=',/,(6e12.4))") colat
!   write(iulog,"('  phifun=',/,(6e12.4))") phifun
!   write(iulog,"('  phifn2=',/,(6e12.4))") phifn2
!   write(iulog,"('  poten =',/,(6e12.4))") poten

    do i=1,nmlon
!    write(iulog,*) ' i, iflag(i) ', i, iflag(i)
      if (iflag(i)==2) then
        poten(i) = (2._r8*(pcen(ihem)-phifun(i))+(phifun(i)-phifn2(i))* &
          0.75_r8)*(colat(i)/theta0(ihem))**3 +                         &
          (1.5_r8*(phifun(i)+phifn2(i))-3._r8*pcen(ihem))*(colat(i)/    &
          theta0(ihem))**2 + 0.75_r8*(phifun(i)-phifn2(i))*(colat(i)/   &
          theta0(ihem)) + pcen(ihem)
      else
        poten(i) = phifun(i)*(max(sin(colat(i)),sinth0)/sinth0)**rr1(ihem)* &
          exp(7._r8*(1._r8-max(sin(colat(i)),sinthr1)/sinthr1))
      endif
    enddo

!   write(iulog,"(/'flwv32: mlat=',i2,' ihem=',i2)") mlat,ihem
!   write(iulog,"('  theta0(ihem)=',e12.4,' pcen(ihem)=',e12.4,' rr1(ihem)=',e12.4)") &
!     theta0(ihem),pcen(ihem),rr1(ihem)
!   write(iulog,"('  sinth0=',e12.4,' sinthr1=',e12.4)") sinth0,sinthr1
!   write(iulog,"('  iflag=',/,(20i3))") iflag
!   write(iulog,"('  dlat =',/,(6e12.4))") dlat
!   write(iulog,"('  dlon =',/,(6e12.4))") dlon
!   write(iulog,"('  colat=',/,(6e12.4))") colat
!   write(iulog,"('  phifun=',/,(6e12.4))") phifun
!   write(iulog,"('  phifn2=',/,(6e12.4))") phifn2
!   write(iulog,"('  poten =',/,(6e12.4))") poten

  end subroutine flwv32
!-----------------------------------------------------------------------
end module heelis