module eos
      implicit none
      include 'physconst.inc'
c
      integer, parameter:: maxel=2
      integer, parameter:: maxion=maxel
c
      integer :: code(maxel)=(/ 100, 200 /) ! element codes
      real*8 :: eps(maxel)=(/ 1.d0, 0.d0 /) ! element abundance
      real*8 :: partial_pressures(0:maxion,maxel)
c
      real*8 :: ion_fraction(0:maxion)
      real*8, private :: T,Pgas,Pe
c
      contains
c
      subroutine zusum(T,Pgas,code,q,chi)
c     -----------------------------------
      implicit none
      real*8, intent(in) :: T,Pgas
      integer, intent(in) :: code
      real*8, intent(out) :: q(0:maxion)
      real*8, intent(out) :: chi(0:maxion)
c
***********************************************************************
* calculate partition function Q and ionization energies chi for 
* (T,Pg) and element code (100*z)
***********************************************************************
c
      q(:) = 0
      chi(:) = 0
      select case(code)
       case(100)
        q(0) = 2.d0
        q(1) = 1.d0
        chi(0) = 13.595d0*ev2erg
       case(200)
        q(0) = 1.d0
        q(1) = 2.d0
        q(2) = 1.d0
        chi(0) = 24.580d0*ev2erg
        chi(1) = 54.403d0*ev2erg
      case default
       write(*,*) 'zusum: code error:',code
       stop 'zusum error'
      end select
      return
c
      end subroutine zusum
c
c
c
      subroutine ionfrac(T,Pgas,Pe,code,ion_fraction,qfak)
c     ----------------------------------------------------
      implicit none
      real*8, intent(in) :: T,Pgas,Pe
      integer, intent(in) :: code
      real*8, intent(out) :: ion_fraction(0:maxion) ! ion fractions
      real*8, intent(out) :: qfak ! total number of charges delivered
c
***********************************************************************
* calculate ion fraction for (T,Pg,Pe) and element code (100*z)
***********************************************************************
c--
c-- local variables:
c--
      real*8 :: q(0:maxion)
      real*8 :: chi(0:maxion)
      real*8 sum,smax,ktinv,tpeup
c--
c-- some useful constants:
c--
      real*8 :: sconst
      real*8 :: const
c
      integer last_ion,ilast,i
c
      save
c
      sconst = (sqrt(2.d0*pi*me*kb)/h)**3*kb
      const = 10.d0**(3.41405d0+1.d0)
      tpeup = 2.5d0*log10(T)-log10(Pe)
      ktinv = 1.d0/(kb*T)
c--
c-- this is the last ion of this element
c--
      last_ion = code/100
      ilast = min(last_ion-1,maxion-1)
c--
c-- compute Q's and load ionization energies
c--
      call zusum(T,Pgas,code,q,chi)
c--
c-- compute saha factors:
c--
      ion_fraction(:) = 0.d0
      do i=0,ilast
       ion_fraction(i+1) = ion_fraction(i)+
     &   log(2.d0*sconst)+log(q(i+1)/q(i))
     &   +(log(10.d0))*tpeup-ktinv*chi(i)
       smax = max(smax,ion_fraction(i+1))
      enddo
c
      do i=0,ilast+1
       ion_fraction(i) = exp(ion_fraction(i)-smax)
      enddo
c
      sum = 0.d0
      do i=ilast+1,0,-1
       sum = sum+ion_fraction(i)
      enddo
c
      sum = 1.d0/sum
      qfak = 0.d0
      do i=0,ilast+1
       qfak = qfak+dble(i)*(ion_fraction(i)*sum)
       ion_fraction(i) = (ion_fraction(i)*sum)
      enddo
c
      return
c
      end subroutine ionfrac
c
c
c
      double precision function pefct(pelog)
c     --------------------------------------
      implicit none
      real*8, intent(in):: pelog
************************************************************************
* compute the iteration function for ions
************************************************************************
c--
c-- local parameter/variables:
c--
      real*8 :: pcharge,qfak
      integer :: i
c
      Pe = 10.d0**pelog
c--
c-- compute the fractions for all elements and their ionization stages:
c--  always compute new partition function to guaranty good results
c--
      pcharge = 0.d0
      do i=1,maxel
       if(code(i) .le. 0) exit 
       call ionfrac(T,Pgas,Pe,code(i),ion_fraction,qfak)
       pcharge = pcharge+qfak*eps(i)*(Pgas-Pe)
      enddo
c
      pefct = pe-pcharge
c
      return
      end function pefct
c
c
c
      double precision function pebrent(func,x1,x2,tol,itused)
c     --------------------------------------------------------
      implicit none
      integer itmax
      real*8 tol,x1,x2,func,eps
      external func
      parameter (itmax = 100,eps=1.d-16)
      integer iter,itused
      real*8 a,b,c,d,e,fa,fb,fc,p,q,r,s,tol1,xm
c
      a = x1
      b = x2
      fa = func(a)
      fb = func(b)
      if((fa .gt. 0.d0 .and. fb .gt. 0.d0) .or. 
     &   (fa .lt. 0.d0 .and. fb .lt. 0.d0)     ) then
       a = log10(1d-300*Pgas)
       b = log10(0.99d0*Pgas)
       fa = func(a)
       fb = func(b)
      endif
c--
c-- for very small T, we may need a much lower initial guess
c--
      if((fa .gt. 0.d0 .and. fb .gt. 0.d0) .or. 
     &   (fa .lt. 0.d0 .and. fb .lt. 0.d0)     ) then
       if(T .lt. 2000.d0) then
        a = log10(1d-50*Pgas)
        b = log10(1d-10*Pgas)
        fa = func(a)
        fb = func(b)
       endif
      endif
c
      if((fa .gt. 0.d0 .and. fb .gt. 0.d0) .or. 
     &   (fa .lt. 0.d0 .and. fb .lt. 0.d0)     ) then
       stop
     &    'pebrent: error in maximum Pe interval detected!'
      endif
      c = b
      fc = fb
      do 11 iter = 1,itmax
        if((fb .gt. 0.d0 .and. fc .gt. 0.d0) .or. 
     &     (fb .lt. 0.d0 .and. fc .lt. 0.d0)) then
          c = a
          fc = fa
          d = b-a
          e = d
        endif
        if(abs(fc) .lt. abs(fb)) then
          a = b
          b = c
          c = a
          fa = fb
          fb = fc
          fc = fa
        endif
        tol1 = 2.d0*eps*abs(b)+0.5d0*tol
        xm = .5*(c-b)
        if(abs(xm) .le. tol1 .or. fb .eq. 0.d0) then
          itused = iter
          x1 = a*(1.d0-1d-3)
          x2 = b*(1.d0+1d-3)
          pebrent = b
          return
        endif
        if(abs(e) .ge. tol1 .and. abs(fa) .gt. abs(fb)) then
          s = fb/fa
          if(a .eq. c) then
            p = 2.d0*xm*s
            q = 1.d0-s
          else
            q = fa/fc
            r = fb/fc
            p = s*(2.d0*xm*q*(q-r)-(b-a)*(r-1.d0))
            q = (q-1.d0)*(r-1.d0)*(s-1.d0)
          endif
          if(p .gt. 0.d0) q = -q
          p = abs(p)
          if(2.d0*p .lt. min(3.d0*xm*q-abs(tol1*q),abs(e*q))) then
            e = d
            d = p/q
          else
            d = xm
            e = d
          endif
        else
          d = xm
          e = d
        endif
        a = b
        fa = fb
        if(abs(d) .gt. tol1) then
          b = b+d
        else
          b = b+sign(tol1,xm)
        endif
        fb = func(b)
11    continue
      write(*,*) 'no convergence in pebrent for ',T,Pgas
      stop 'pebrent exceeding maximum iterations'
      end function pebrent
c
c
c
      subroutine eos_ions(t_in,pg_in,pe_out,itused)
c     ---------------------------------------------
      implicit none
      real*8, intent(in) :: t_in,pg_in
      real*8, intent(inout) :: pe_out
      integer, intent(out) :: itused
************************************************************************
* solution of the EOS for ions and atoms using Brent's method
*
*-- input:
* t_in: temperature (k)
* pg_in: gas pressure (dyn/cm**2)
* pe_out: electron pressure (dyn/cm**2)
* mode: slot to store the results into
*-- output:
* pe: electron pressure (dyn/cm**2)
* itused: number of iterations used
*
************************************************************************
c
      real*8 pnuc,pemin,pemax,pgold,hlp,pelog,qfak
      integer i,j
      logical, save :: first=.true.
c--
c-- init's:
c--
      t = t_in
      pgas = pg_in
c--
c-- get inital guesses for Pe:
c--
      if(first) then
       pemin = log10(1d-10*pgas)
       pemax = log10(0.70*pgas)
       first = .false.
       pgold = log10(pgas)
      else
       hlp = log10(pgas)
       pemin = pemin + hlp - pgold
       pemax = pemax + hlp - pgold
       pgold = hlp
      endif
c--
c-- check if the interval works, correct if not
c--
c--
c-- use brent to solve for Pe:
c--
      pelog = pebrent(pefct,pemin,pemax,1d-8,itused)
      pe_out = 10.d0**pelog
c--
c-- distribute the results to target array:
c--
      pnuc = pgas-pe
      partial_pressures(:,:) = 0
      do i=1,maxel
       call ionfrac(T,Pgas,Pe,code(i),ion_fraction,qfak)
       do j=0,code(i)/100
        partial_pressures(j,i) = ion_fraction(j)*eps(i)*pnuc
       enddo
      enddo
c
      return
      end subroutine eos_ions
c
      end module eos
c
c
c
      program eos_table
      use eos
      implicit none
c
      real*8 T,Pgas,Pe,qfak
      real*8 Pea,PH1,PH2
      integer :: i,itused
c
      real*8 :: q(0:maxion),chi(0:maxion)
c
      code(1) = 100 ! H
      code(2) = 200 ! He
c
      eps(1) = 1.d0 ! eps(H)
      eps(2) = 0.d0 ! eps(He)
c
      eps(:) = eps(:)/sum(eps(:)) ! re-normalize
c
      T = 5000.d0
      Pgas = 100.d0
      Pe = Pgas*0.5d0
c--
c-- jump to different tests:
c--
      goto 300
c--
c-- test: Q's:
c--
  100 continue
      call zusum(T,Pgas,100,q,chi)
      write(*,*) 'Q:',100,q(:),chi(:)
      call zusum(T,Pgas,200,q,chi)
      write(*,*) 'Q:',200,q(:),chi(:)
      stop 'Q test done'
c
  200 continue
      call ionfrac(T,Pgas,Pe,100,ion_fraction,qfak)
      write(*,*) 'ionfrac:',100,qfak
      write(*,*) 'ionfrac:',100,ion_fraction(:)
      call ionfrac(T,Pgas,Pe,200,ion_fraction,qfak)
      write(*,*) 'ionfrac:',200,qfak
      write(*,*) 'ionfrac:',200,ion_fraction(:)
      stop 'ionfrac test done'
c--
c-- test the real thing:
c--
  300 continue
      write(*,*) 'enter T,Pgas:'
      read(*,*) T,Pgas
      call eos_ions(t,pgas,pe,itused)
      write(*,*) 'results for (T,pgas):',T,Pgas
      write(*,*) 'Pe=',Pe
      write(*,*) 'iterations:',itused
      write(*,*) 'partial pressures'
      do i=1,maxel
       write(*,*) 'code=',code(i)
       write(*,*) partial_pressures(0:code(i)/100,i)
      enddo
c--
c-- compare to analytic solution for Pe:
c--
      call analytic(T,Pgas,Pea,PH1,PH2)
      write(*,*) "analytic Pe =",Pea," error=",(Pe-Pea)/Pe
      write(*,*) "analytic PH1=",PH1," error=",
     &           (partial_pressures(0,1)-PH1)/PH1
      write(*,*) "analytic PH2=",PH2," error=",
     &           (partial_pressures(1,1)-PH2)/PH2
      goto 300
c
      contains
      subroutine analytic(T,Pg,Pe,PH1,PH2)
      implicit none
      real*8 :: T,Pg,Pe,PH1,PH2
c
      include 'physconst.inc'
      real*8 ::N,Ne,phi,q1,q2,chi
c
      N = Pg/(kb*T)
c
      q1 = 2.d0
      q2 = 1.d0
      chi = 13.595d0*ev2erg
      phi = 2.d0*q2/q1*(2.d0*pi*me*kb*T/h**2)**1.5d0*exp(-chi/(kb*T))
c
      Ne = phi*(sqrt(1.d0+N/phi)-1.d0)
c
      Pe = Ne*kb*T
      PH2 = Pe
      PH1 = Pg-2*Pe
c
      return
      end subroutine analytic
      end