#include "cctk.h"
#include "cctk_Parameters.h"
#include "cctk_Arguments.h"
#include "cctk_Functions.h"

! eoskey:
! 1 --- polytropic EOS
! 2 --- gamma-law EOS
! 3 --- hybrid EOS
! 4 --- finite-T microphysical NSE EOS

subroutine EOS_Omni_EOS_Press(eoskey,keytemp,rf_precision,npoints,&
                              rho,eps,temp,ye,press,keyerr,anyerr)

  use EOS_Omni_Module
  implicit none
  DECLARE_CCTK_PARAMETERS

  CCTK_INT, intent(in)     :: eoskey,keytemp,npoints
  CCTK_INT, intent(out)    :: keyerr(npoints)
  CCTK_INT, intent(out)    :: anyerr
  CCTK_REAL, intent(in)    :: rf_precision
  CCTK_REAL, intent(in)    :: rho(npoints),ye(npoints)
  CCTK_REAL, intent(inout) :: eps(npoints), temp(npoints)
  CCTK_REAL, intent(out)   :: press(npoints)

  ! local vars
  integer          :: i
  character(256)   :: warnstring
  real*8           :: hybrid_local_gamma, hybrid_local_k_cgs, &
                      hybrid_p_poly, hybrid_p_th
  real*8,parameter :: zero = 0.0d0
  ! temporary vars for nuc_eos
  real*8           :: xrho,xye,xtemp,xenr,xent
  real*8           :: xprs,xmunu,xcs2
  real*8           :: xdedt,xdpderho,xdpdrhoe

  anyerr    = 0
  keyerr(:) = 0

  select case (eoskey)
     case (1)
        ! polytropic EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * poly_k_cgs * &
                   (rho(i)*inv_rho_gf)**(poly_gamma) / &
                   (poly_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           press(i) = press_gf * poly_k_cgs * &
                (rho(i)*inv_rho_gf)**poly_gamma
        enddo
     case (2)
        ! gamma-law EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * gl_k_cgs * &
                   (rho(i)*inv_rho_gf)**(gl_gamma) / &
                   (gl_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           press(i) = (gl_gamma - 1.0d0) * rho(i) * eps(i)
        enddo

     case (3)
        ! hybrid EOS
        do i=1,npoints
           if(rho(i).gt.hybrid_rho_nuc) then
              hybrid_local_gamma = hybrid_gamma2
              hybrid_local_k_cgs = hybrid_k2_cgs
           else
              hybrid_local_gamma = hybrid_gamma1
              hybrid_local_k_cgs = hybrid_k1_cgs
           endif
           hybrid_p_poly = press_gf * hybrid_local_k_cgs * &
                     (rho(i) * inv_rho_gf)**hybrid_local_gamma
           hybrid_p_th = - press_gf * hybrid_local_k_cgs * (hybrid_gamma_th - 1.d0) /      &
                (hybrid_local_gamma - 1.0d0) * (rho(i) * inv_rho_gf)**hybrid_local_gamma + &
                (hybrid_gamma_th - 1.0d0) * rho(i) * eps(i) -                              &
                (hybrid_gamma_th - 1.d0) * (hybrid_local_gamma - hybrid_gamma1) /          &
                (hybrid_gamma1 - 1.d0) / (hybrid_gamma2 - 1.d0) *                          &
                press_gf * hybrid_k1_cgs * inv_rho_gf**hybrid_gamma1 *                     &
                hybrid_rho_nuc**(hybrid_gamma1 - 1.d0) * rho(i)
           hybrid_p_th = max(zero, hybrid_p_th)
           press(i) = hybrid_p_poly + hybrid_p_th
        enddo

     case (4)

        do i=1,npoints

           xrho = rho(i) * inv_rho_gf
           xtemp = temp(i)
           xye = ye(i)
           xenr = eps(i) * inv_eps_gf
           call nuc_eos_press_eps(xrho,xtemp,xye,xenr,xprs,& 
                keytemp,keyerr(i),rf_precision)
           
           if(keyerr(i).ne.0) then
              anyerr = 1
           endif

           if(keytemp.eq.1) then
              eps(i) = xenr * eps_gf
           else
              temp(i) = xtemp
           endif

           press(i) = xprs * press_gf

        enddo

     case DEFAULT
        write(warnstring,*) "eoskey ",eoskey," not implemented!"
        call CCTK_WARN(0,warnstring)
     end select

end subroutine EOS_Omni_EOS_Press



subroutine EOS_Omni_EOS_DPressByDRho(eoskey,keytemp,rf_precision,npoints,&
                              rho,eps,temp,ye,dpdrhoe,keyerr,anyerr)

  use EOS_Omni_Module
  implicit none
  DECLARE_CCTK_PARAMETERS

  CCTK_INT, intent(in)     :: eoskey,keytemp,npoints
  CCTK_INT, intent(out)    :: keyerr(npoints)
  CCTK_INT, intent(out)    :: anyerr
  CCTK_REAL, intent(in)    :: rf_precision
  CCTK_REAL, intent(in)    :: rho(npoints),ye(npoints)
  CCTK_REAL, intent(inout) :: eps(npoints), temp(npoints)
  CCTK_REAL, intent(out)   :: dpdrhoe(npoints)

  ! local vars
  integer          :: i
  character(256)   :: warnstring
  real*8           :: hybrid_local_gamma, hybrid_local_k_cgs, &
                      hybrid_dp_poly, hybrid_dp_th1, hybrid_dp_th2
  real*8,parameter :: zero = 0.0d0
  ! temporary vars for nuc_eos
  real*8           :: xrho,xye,xtemp,xenr,xent
  real*8           :: xprs,xmunu,xcs2
  real*8           :: xdedt,xdpderho,xdpdrhoe

  anyerr    = 0
  keyerr(:) = 0

  select case (eoskey)
     case (1)
        ! polytropic EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * poly_k_cgs * &
                   (rho(i)*inv_rho_gf)**(poly_gamma) / &
                   (poly_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           dpdrhoe(i) = press_gf * poly_k_cgs *  &
                poly_gamma * inv_rho_gf *        & 
                (rho(i)*inv_rho_gf) ** (poly_gamma - 1.d0)
        enddo
     case (2)
        ! gamma-law EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * gl_k_cgs * &
                   (rho(i)*inv_rho_gf)**(gl_gamma) / &
                   (gl_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           dpdrhoe(i) = (gl_gamma-1.0d0) * &
                eps(i)
        enddo
     case (3)
        ! hybrid EOS
        do i=1,npoints
           if(rho(i).gt.hybrid_rho_nuc) then
              hybrid_local_gamma = hybrid_gamma2
              hybrid_local_k_cgs = hybrid_k2_cgs
           else
              hybrid_local_gamma = hybrid_gamma1
              hybrid_local_k_cgs = hybrid_k1_cgs
           endif
           hybrid_dp_poly = hybrid_local_gamma * press_gf *                 &
                hybrid_local_k_cgs * rho(i)**(hybrid_local_gamma - 1.0d0) * &
                inv_rho_gf**hybrid_local_gamma

           hybrid_dp_th1 = - hybrid_local_gamma * press_gf * hybrid_local_k_cgs * &
                (hybrid_gamma_th - 1.d0) / (hybrid_local_gamma - 1.d0) *      &
                rho(i)**(hybrid_local_gamma - 1.d0) * inv_rho_gf**hybrid_local_gamma

           hybrid_dp_th2  = (hybrid_gamma_th - 1.d0) * eps(i)                        &
                - (hybrid_gamma_th - 1.d0) * (hybrid_local_gamma - hybrid_gamma1) /  &
                (hybrid_gamma1 - 1.d0) / (hybrid_gamma2 - 1.d0) *                    &
                press_gf * hybrid_k1_cgs * inv_rho_gf**hybrid_gamma1 *               &
                hybrid_rho_nuc**(hybrid_gamma1 - 1.d0)

           dpdrhoe(i) = hybrid_dp_poly + hybrid_dp_th1 + hybrid_dp_th2

        enddo

     case (4)
        do i=1,npoints
           xrho = rho(i) * inv_rho_gf
           xtemp = temp(i)
           xye = ye(i)
           xenr = eps(i) * inv_eps_gf
           call nuc_eos_short(xrho,xtemp,xye,xenr,xprs,& 
                xent,xcs2,xdedt,xdpderho,xdpdrhoe,xmunu,&
                keytemp,keyerr(i),rf_precision)
           
           if(keyerr(i).ne.0) then
              anyerr = 1
           endif

           if(keytemp.eq.1) then
              eps(i) = xenr * eps_gf
           else
              temp(i) = xtemp
           endif

           dpdrhoe(i) = xdpdrhoe * press_gf * inv_rho_gf
        enddo
     case DEFAULT
        write(warnstring,*) "eoskey ",eoskey," not implemented!"
        call CCTK_WARN(0,warnstring)
   end select
     
end subroutine EOS_Omni_EOS_DPressByDRho

subroutine EOS_Omni_EOS_DPressByDEps(eoskey,keytemp,rf_precision,npoints,&
                              rho,eps,temp,ye,dpdepsrho,keyerr,anyerr)

  use EOS_Omni_Module
  implicit none
  DECLARE_CCTK_PARAMETERS

  CCTK_INT, intent(in)     :: eoskey,keytemp,npoints
  CCTK_INT, intent(out)    :: keyerr(npoints)
  CCTK_INT, intent(out)    :: anyerr
  CCTK_REAL, intent(in)    :: rf_precision
  CCTK_REAL, intent(in)    :: rho(npoints),ye(npoints)
  CCTK_REAL, intent(inout) :: eps(npoints), temp(npoints)
  CCTK_REAL, intent(out)   :: dpdepsrho(npoints)

  ! local vars
  integer          :: i
  character(256)   :: warnstring
  ! temporary vars for nuc_eos
  real*8           :: xrho,xye,xtemp,xenr,xent
  real*8           :: xprs,xmunu,xcs2
  real*8           :: xdedt,xdpderho,xdpdrhoe

  anyerr    = 0
  keyerr(:) = 0

  select case (eoskey)
     case (1)
        ! polytropic EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * poly_k_cgs * &
                   (rho(i)*inv_rho_gf)**(poly_gamma) / &
                   (poly_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           dpdepsrho(i) = 0.0d0
        enddo
     case (2)
        ! gamma-law EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * gl_k_cgs * &
                   (rho(i)*inv_rho_gf)**(gl_gamma) / &
                   (gl_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           dpdepsrho(i) = (gl_gamma - 1.0d0) * &
                rho(i)
        enddo
     case (3)
        ! hybrid EOS
        do i=1,npoints
           dpdepsrho(i) = (hybrid_gamma_th - 1.0d0) * rho(i)
        enddo

     case (4)
        ! nuc_eos
        do i=1,npoints

           xrho = rho(i) * inv_rho_gf
           xtemp = temp(i)
           xye = ye(i)
           xenr = eps(i) * inv_eps_gf
           call nuc_eos_short(xrho,xtemp,xye,xenr,xprs,& 
                xent,xcs2,xdedt,xdpderho,xdpdrhoe,xmunu,&
                keytemp,keyerr(i),rf_precision)
           
           if(keyerr(i).ne.0) then
              anyerr = 1
           endif

           if(keytemp.eq.1) then
              eps(i) = xenr * eps_gf
           else
              temp(i) = xtemp
           endif

           dpdepsrho(i) = xdpderho * press_gf * inv_eps_gf 

        enddo

     case DEFAULT
        write(warnstring,*) "eoskey ",eoskey," not implemented!"
        call CCTK_WARN(0,warnstring)
   end select
     
end subroutine EOS_Omni_EOS_DPressByDEps

subroutine EOS_Omni_EOS_cs2(eoskey,keytemp,rf_precision,npoints,&
                            rho,eps,temp,ye,cs2,keyerr,anyerr)

  use EOS_Omni_Module
  implicit none
  DECLARE_CCTK_PARAMETERS

  CCTK_INT, intent(in)     :: eoskey,keytemp,npoints
  CCTK_INT, intent(out)    :: keyerr(npoints)
  CCTK_INT, intent(out)    :: anyerr
  CCTK_REAL, intent(in)    :: rf_precision
  CCTK_REAL, intent(in)    :: rho(npoints),ye(npoints)
  CCTK_REAL, intent(inout) :: eps(npoints), temp(npoints)
  CCTK_REAL, intent(out)   :: cs2(npoints)

  ! local vars
  integer          :: i
  character(256)   :: warnstring
  real*8           :: xpress,xdpdrhoe,xdpderho
  real*8           :: hybrid_local_gamma, hybrid_local_k_cgs, &
                      hybrid_p_poly, hybrid_p_th
  real*8,parameter :: zero = 0.0d0
  ! temporary vars for nuc_eos
  real*8           :: xrho,xye,xtemp,xenr,xent
  real*8           :: xprs,xmunu,xcs2
  real*8           :: xdedt

  anyerr    = 0
  keyerr(:) = 0

  select case (eoskey)
     case (1)
        ! polytropic EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * poly_k_cgs * &
                   (rho(i)*inv_rho_gf)**(poly_gamma) / &
                   (poly_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           xpress = press_gf*poly_k_cgs * &
                (rho(i)*inv_rho_gf)**(poly_gamma)
           cs2(i) = poly_gamma * xpress / rho(i) / &
                (1 + eps(i) + xpress/rho(i))
        enddo
     case (2)
        ! gamma-law EOS
        if(keytemp.eq.1) then
           do i=1,npoints
              eps(i) = press_gf * gl_k_cgs * &
                   (rho(i)*inv_rho_gf)**(gl_gamma) / &
                   (gl_gamma - 1.0d0) / rho(i)
           enddo
        endif
        do i=1,npoints
           xpress = (gl_gamma-1.0d0)*rho(i)*eps(i)
           xdpdrhoe = (gl_gamma-1.0d0)*eps(i)
           xdpderho = (gl_gamma-1.0d0)*rho(i)
           cs2(i) = (xdpdrhoe + xpress * xdpderho / (rho(i)**2)) / &
                (1.0d0 + eps(i) + xpress/rho(i))
        enddo
     case(3)
        ! hybrid EOS
        do i=1,npoints
           if(rho(i).gt.hybrid_rho_nuc) then
              hybrid_local_gamma = hybrid_gamma2
              hybrid_local_k_cgs = hybrid_k2_cgs
                 else
              hybrid_local_gamma = hybrid_gamma1
              hybrid_local_k_cgs = hybrid_k1_cgs
           endif
           ! first calculate the pressure
           hybrid_p_poly = press_gf * hybrid_local_k_cgs * &
                     (rho(i) * inv_rho_gf)**hybrid_local_gamma
           hybrid_p_th = - press_gf * hybrid_local_k_cgs * (hybrid_gamma_th - 1.d0) /      &
                (hybrid_local_gamma - 1.0d0) * (rho(i) * inv_rho_gf)**hybrid_local_gamma + &
                (hybrid_gamma_th - 1.0d0) * rho(i) * eps(i) -                              &
                (hybrid_gamma_th - 1.d0) * (hybrid_local_gamma - hybrid_gamma1) /          &
                (hybrid_gamma1 - 1.d0) / (hybrid_gamma2 - 1.d0) *                          &
                press_gf * hybrid_k1_cgs * inv_rho_gf**hybrid_gamma1 *                     &
                hybrid_rho_nuc**(hybrid_gamma1 - 1.d0) * rho(i)
           hybrid_p_th = max(zero, hybrid_p_th)
           xpress = hybrid_p_poly + hybrid_p_th 
           cs2(i) = (hybrid_local_gamma * hybrid_p_poly + hybrid_gamma_th * hybrid_p_th) / &
                rho(i) / (1.0d0 + eps(i) + xpress/rho(i))
        enddo
     case(4)
        ! nuc_eos

        do i=1,npoints
           
           xrho = rho(i) * inv_rho_gf
           xtemp = temp(i)
           xye = ye(i)
           xenr = eps(i) * inv_eps_gf
           call nuc_eos_short(xrho,xtemp,xye,xenr,xprs,& 
                xent,xcs2,xdedt,xdpderho,xdpdrhoe,xmunu,&
                keytemp,keyerr(i),rf_precision)
           
           if(keyerr(i).ne.0) then
              anyerr = 1
           endif

           if(keytemp.eq.1) then
              eps(i) = xenr * eps_gf
           else
              temp(i) = xtemp
           endif

           cs2(i) = xcs2 * cliteinv2 / &
                (1.0d0 + eps(i) + xprs * press_gf / rho(i))

        enddo


     case DEFAULT
        write(warnstring,*) "eoskey ",eoskey," not implemented!"
        call CCTK_WARN(0,warnstring)
   end select
     
end subroutine EOS_Omni_EOS_cs2


subroutine EOS_Omni_EOS_eps_from_press(eoskey,keytemp,rf_precision,npoints,&
                              rho,eps,temp,ye,press,xeps,keyerr,anyerr)

  use EOS_Omni_Module
  implicit none
  DECLARE_CCTK_PARAMETERS

  CCTK_INT, intent(in)     :: eoskey,keytemp,npoints
  CCTK_INT, intent(out)    :: keyerr(npoints)
  CCTK_INT, intent(out)    :: anyerr
  CCTK_REAL, intent(in)    :: rf_precision
  CCTK_REAL, intent(in)    :: rho(npoints),ye(npoints),press(npoints)
  CCTK_REAL, intent(inout) :: eps(npoints), temp(npoints)
  CCTK_REAL, intent(out)   :: xeps(npoints)

  ! local vars
  integer        :: i
  character(256) :: warnstring
  ! temporary vars for nuc_eos
  real*8           :: xrho,xye,xtemp,xenr,xent
  real*8           :: xprs,xmunu,xcs2
  real*8           :: xdedt,xdpderho,xdpdrhoe

  if(keytemp.eq.1) then
     anyerr = 1
     keyerr(:) = -1
  else
     anyerr    = 0
     keyerr(:) = 0
  endif

  select case (eoskey)
     case (1)
        ! polytropic EOS
        do i=1,npoints
           xeps(i) = press(i) / (poly_gamma - 1.0d0) / rho(i)
        enddo
     case (2)
        ! gamma-law EOS
        do i=1,npoints
           xeps(i) = press(i) / (gl_gamma - 1.0d0) / rho(i)
        enddo
     case (3)
        ! hybrid EOS
        write(warnstring,*) "EOS_Omni_EpsFromPress call not supported for hybrid EOS"
        call CCTK_WARN(0,warnstring)
     case (4)
        ! nuc EOS
        write(warnstring,*) "EOS_Omni_EpsFromPress call not supported for nuc_eos yet"
        call CCTK_WARN(0,warnstring)
     case DEFAULT
        write(warnstring,*) "eoskey ",eoskey," not implemented!"
        call CCTK_WARN(0,warnstring)
     end select

end subroutine EOS_Omni_EOS_eps_from_press


subroutine EOS_Omni_EOS_RhoFromPressEpsTempEnt(eoskey,&
                           ikeytemp,rf_precision,&
                           npoints,rho,eps,temp,entropy,ye,press,keyerr,anyerr)

  use EOS_Omni_Module
  implicit none
  DECLARE_CCTK_PARAMETERS

  CCTK_INT, intent(in)     :: eoskey,ikeytemp,npoints
  CCTK_INT, intent(out)    :: keyerr(npoints)
  CCTK_INT, intent(out)    :: anyerr
  CCTK_REAL, intent(in)    :: rf_precision
  CCTK_REAL, intent(in)    :: ye(npoints),press(npoints)
  CCTK_REAL, intent(inout) :: rho(npoints),temp(npoints),eps(npoints)
  CCTK_REAL, intent(inout) :: entropy(npoints)


  ! local vars
  integer        :: i
  character(256) :: warnstring
  integer :: keytemp
  ! temporary vars for nuc_eos
  real*8           :: xrho,xye,xtemp,xenr,xent
  real*8           :: xprs,xmunu,xcs2
  real*8           :: xdedt,xdpderho,xdpdrhoe
  
  keytemp = ikeytemp

  anyerr    = 0
  keyerr(:) = 0

  select case (eoskey)
     case (1)
        ! polytropic EOS
        do i=1,npoints
           rho(i) = (press(i) / poly_k)**(1.0/poly_gamma)
           eps(i) = press(i) / (poly_gamma - 1.0d0) / rho(i)
        enddo
     case (2)
        ! gamma-law EOS
        do i=1,npoints
           rho(i) = press(i) / ((gl_gamma - 1.0d0)*eps(i))
        enddo
     case (3)
        ! hybrid EOS
        write(warnstring,*) "EOS_Omni_RestMassDensityFromPressEpsTemp not supported for hybrid EOS"
        call CCTK_WARN(0,warnstring)
     case (4)
        ! nuc EOS
        if(ikeytemp.eq.2) then
           call CCTK_WARN(0,"This function does not work yet when coming in with entropy")
        else if(ikeytemp.eq.1) then
           keytemp = 3
        else
           call CCTK_WARN(0,"This function does not work yet when coming in with this keytemp")
        endif

        do i=1,npoints
           ! initial guess
           xprs = press(i) * inv_press_gf
           xrho = rho(i) * inv_rho_gf
           xtemp = temp(i)
           xent = entropy(i)
           xye = ye(i)
           xenr = eps(i) * inv_eps_gf
           call nuc_eos_short(xrho,xtemp,xye,xenr,xprs,& 
                xent,xcs2,xdedt,xdpderho,xdpdrhoe,xmunu,&
                keytemp,keyerr(i),rf_precision)
           
           if(keyerr(i).ne.0) then
              anyerr = 1
           endif

           if(keytemp.eq.3) then
              eps(i) = xenr * eps_gf
              rho(i) = xrho * rho_gf
              entropy(i) = xent
           else
              call CCTK_WARN(0, "This keytemp is not implemented yet")
           endif

        enddo
     case DEFAULT
        write(warnstring,*) "eoskey ",eoskey," not implemented!"
        call CCTK_WARN(0,warnstring)
     end select

end subroutine EOS_Omni_EOS_RhoFromPressEpsTempEnt