from ClimateUtilities import *
import phys,planets
import math

#Extra monatomic gases not in phys
AtomicH = Dummy()
AtomicH.R = phys.Rstar
AtomicH.Rcp = 1.- 3./5.
AtomicH.MolecularWeight = 1.
AtomicH.cp = AtomicH.R/AtomicH.Rcp

AtomicO = Dummy()
AtomicO.R = phys.Rstar/16.
AtomicO.Rcp = 1.- 3./5.
AtomicO.MolecularWeight = 16.
AtomicO.cp = AtomicO.R/AtomicO.Rcp

#Fictitious gas with diatomic R/cp but arbitrary
#molecular weight (useful for exploring how things scale
#with the molecular weight)

Oobleck = Dummy()
Oobleck.MolecularWeight = 5.
Oobleck.R = phys.Rstar/Oobleck.MolecularWeight
Oobleck.Rcp = phys.H2.Rcp
Oobleck.cp = Oobleck.R/Oobleck.Rcp

#-----------------------------------------------------------------
#Query: When Q0 is too large, T gets driven to zero as we integrate
#backward. What is the implied T at the lower boundary, using
#energy conservation?  How does the solution fail when we integrate
#forward from there?  (Look at this analytically, assuming Mach
#number is small at the boundary.  Graphically, it looks like
#the temperature should be driven negative (due to adiabatic expansion)
#integrating forward from lower boundary, but this is the opposite
#of what you would expect in the case of large heating. But note:
#This question is ill-posed, since in the case of heating we don't know
#theta at the lower boundary if we can't integrate the entropy equation.
#through to there.  Thus, the solution failure is DIFFERENT from the
#solution failure in the adiabatic case with too low temperature, where
#we run out of energy. Basically, when T gets driven to zero due to too
#large heating, we just have to increase the mass flux until we can
#get through the temperature minimum. This can be looked at by
#using the equation for the location and value of the minimum temperature
#(easy if Mach number is small near the temperature minimum).
#
#Note:  Solution just depends on F/Q0 . If we know the solution
#for one F/Q0 (for a given rcrit), just rescale n0 by Q0 to get other sol'ns
#Hence, for a given Q0, it's really just rcrit that must be varied to
#determine the family of solutions.
#----------------------------------------------------------------

#The radius r passed to F is nondimensional, scaled
#by r0.  M = rho0*w0;  total mass flux is this times 4 pi r0**2
#Note that 2*g0*r0 is square of escape velocity.

#---------Routines for making energy conservation plots-----------
#These plots are only for the no-heating case. If we want to illustrate
#the intersection geometry including heating, we need to make
#theta change with r, and also put in the proper evaluation of E0(r)
#allowing for heating.
#
#Note that n0 changes the value of F, but not the shape of the graph
#Graphs just depend on F/rho0, which is c0*M0*r0**2
def EMach1(M,params):
    params.M = M
    cSquared = gam*R*T1(params)
    return cSquared*(M*M/2. + beta/(1.-beta)) - gs*rs/params.r - params.E0
def MachPlotCrit(rcrit,rL):
    #Initialize at the critical point
    c0 = (.5*gs*rs/rcrit)**.5
    T0 = c0*c0*beta/R
    n0 = 1. #Value  of n0 doesn't matter for this. It cancels out.
    params = Init(1.,n0,T0,rcrit)
    params.E0 = 0. # Don't offset energy
    #params.E0 = params.E00
    return MachPlot(params,rL)
def MachPlot1(M0,r0,T0,rL):
    n0 = 1. # Value of n0 doesn't matter for this
    params = Init(M0,n0,T0,r0)
    params.E0 = 0.
    #params.E0 = params.E00
    return MachPlot(params,rL)
def MachPlot(params,rList):
    c = Curve()
    c.addCurve(MachList)
    for rr in rList:
        params.r = rr
        Elist = [EMach1(M,params) for M in MachList]
        c.addCurve(Elist,'E%f'%rr)
    c.XlogAxis  = True
    return c
def T1(params):
    ctheta = (R*params.theta/beta)**.5
    rhotheta = params.p0/(R*params.theta)
    b = 2.*(1.-beta)/(1.+beta)
    return params.theta*( params.F/(params.r**2 *params.M*rhotheta*ctheta))**b
#-------------------------------------------------------------------


#Routines for integrating solution including heating
def EMach(M,params):
    params.M = M
    cSquared = gam*R*T(params)
    return cSquared*(M*M/2. + beta/(1.-beta)) - gs*rs/params.r - params.E0
#The small M approximation here only works when T is an equilibrium
#solution of E(M) = 0. For the graphical illustrations, use T1, which
#doesn't have the small M approximation.
def T(params):
    ctheta = (R*params.theta/beta)**.5
    rhotheta = params.p0/(R*params.theta)
    b = 2.*(1.-beta)/(1.+beta)
    if params.M > .01:
        return params.theta*( params.F/(params.r**2 *params.M*rhotheta*ctheta))**b
    else:
        TT = (gs*rs/params.r + params.E0)/(R*gam*(beta/(1.-beta)))
        return TT


#This initialization should work for any point, not just
#the trans-sonic point. Note that r0 is in units of planetary radius
#**WARNING: This gives the wrong initialization if M <= .01,
#  since it uses EMach, which solves the energy equation for T
#  directly when M is small (i.e. hydrostatic balance). 
#  Need to fix this.  
def Init(M0,n0,T0,r0):
    c0 = (gam*R*T0)**.5
    rho0 = n0*m
    F = c0*M0*rho0*r0**2
    p0 = rho0*R*T0
    #
    params = Dummy()
    params.E0 = 0.
    params.r = r0
    params.F = F
    params.p0 = p0
    params.theta = T0
    params.E00 = EMach(M0,params) - r0*r0*Frad(r0)/F
    return params
    
    
#Radiative flux functions
def Frad(r):
    if r > rlim:
        return Q0*rh3*(1.-math.exp(-(rlim-rh1)/rh3))/(r*r)
    elif (r>rh1) and (r <= rlim):
        return Q0*rh3*(1.-math.exp(-(r-rh1)/rh3))/(r*r)
    else:
        return 0.

def Qrad(r):
    if r > rlim:
        return 0.
    elif (r>rh1) and (r <= rlim):
        return Q0*math.exp(-(r-rh1)/rh3)
    else:
        return 0.
    
#Set up Newton solver for getting Mach number
ans = newtSolve(EMach)
#Derivative function for evaluating theta
def fTheta(r,logtheta,params):
    params.r = r
    params.E0 =  params.E00+r*r*Frad(r)/(params.F)
    ans.setParams(params)
    params.theta = math.exp(logtheta)
    if (params.M < .01):
        #Use analytic solution for M and T
        TT = T(params)
        p = params.p0*(TT/params.theta)**(1./Rcp)
        rho = p/(R*TT)
        c = (gam*R*TT)**.5
        params.M = (params.F/(rho*c*params.r**2))
    else:
        try:
            params.M = ans(params.M)
            if params.M == 'No Convergence':
                print 'Convergence failure in fTheta',params.r,params.theta
        except:
            print 'Convergence failure in fTheta',params.r,params.theta
        TT = T(params)
    return Qrad(r)/(cp*TT*params.F) #Qrad is heating per unit depth

def Solve(r,params):
    params.r = r
    params.E0 =  params.E00+r*r*Frad(r)/(params.F)
    ans.setParams(params)
    if (params.M < .01):
        #Use analytic solution for M and T
        TT = T(params)
        p = params.p0*(TT/params.theta)**(1./Rcp)
        rho = p/(R*TT)
        c = gam*R*TT
        params.M = (params.F/(rho*c*params.r**2))
    else:
        try:
            params.M = ans(params.M)
            if params.M == 'No Convergence':
                print 'Convergence failure in fTheta',params.r,params.theta
        except:
            print 'Convergence failure in fTheta',params.r,params.theta
        TT = T(params)
    return TT,params.M

#Set parameters for transonic point
#**ToDo: Will need to fix up energy plot and MachrPlot routines
#

#Specify the two controlling parameters. These determine the
#outflow mass flux F. They can be adjusted to meet the
#boundary conditions on temperature and particle density at
#the lower boundary r
#


#We start at the critical point, and then integrate inward
#Later,F and rcrit (or n0 and rcrit) should be made arguments and
#p0 etc. should
#be computed. To do that we will need to change some existing
#globals. We will need a different version to do subcritical cases
#The following calculation initializes at the critical point
#Try:
#        rcrit = 35.
#        n0 = 1.e9
def doCalc(n0,rcrit,switch = 'Backward',mode = 'Plot'):
    # Initialization (Assumes heating is zero at the trans-sonic point)
    c0 = (.5*(gs/rcrit**2)*rs*rcrit)**.5
    T0 = c0*c0*beta/R
    params = Init(1.,n0,T0,rcrit)
    print 'Energy',params.E00
    if switch == 'Backward':
        rstart = rcrit - .01
        params.M = .95
        dr = -.01
    else:
        rstart = rcrit + .01
        params.M = 1.05
        dr = .01
    #
    #Create lists for saving output
    rL = []
    ML = []
    TL = []
    thetaL = []
    nL = []
    #
    # Carry out the integration
    #
    ThetaInt = integrator(fTheta,rstart,math.log(T0),dr)
    ThetaInt.setParams(params)
    while (True):
        if dr < 0.:
            if abs(params.r - 1.1)< .0001:
                break
        else:
            if params.r > 100.:
                break
        try:
            values = ThetaInt.next(dr)
        except:
            break
        if params.M == 'No Convergence':
            break
        params.theta = math.exp(values[1])
        params.r = values[0]
        rL.append(params.r)
        ML.append(params.M)
        thetaL.append(params.theta)
        TL.append(T(params))
        #Density
        p = params.p0*(TL[-1]/params.theta)**(1./Rcp)
        nL.append((p/(R*TL[-1]))/m)
#
#   Print out terminal results here
    p = params.p0*(TL[-1]/params.theta)**(1./Rcp)
    nbot = (p/(R*TL[-1]))/m
    print 'rmin,pbot,Tbot,nbot,flux = ',rL[-1],p/1.e5,TL[-1],nbot,params.F
#
#   Return results for plotting
#
    c = Curve()
    c.addCurve(rL)
    c.addCurve(ML)

    c1 = Curve()
    c1.addCurve(rL)
    c1.addCurve(TL)
    c1.addCurve(thetaL)

    c2 = Curve()
    c2.addCurve(rL)
    c2.addCurve(nL)
    
    if mode == 'Plot':
        return c,c1,c2
    else:
        return TL[-1] - Tbot
    return TL[-1] - Tbot

#Identical to doCalc, except it starts from the bottom.
#Note the special initialization used
def doCalc1(n0,T0,F,rmax = 300.,rb=1.1):
    # Initialization at base
    #  Note that this initialization assumes the radiative flux
    #  vanishes at the base
    params = Dummy()
    params.r = rb
    params.F = F
    rho0 = n0*m
    params.p0 = rho0*R*T0
    params.theta = T0
    w = F/(rho0*rb**2)
    params.E00 = .5*w*w + cp*T0 - gs*rs/rb
    params.E0 = .5*w*w + cp*T0 - gs*rs/rb  
    print 'Energy',params.E00
#
    rstart = rb
    params.M = .001
    dr = .01
#
    #
    #Create lists for saving output
    rL = []
    ML = []
    TL = []
    thetaL = []
    nL = []
    #
    # Carry out the integration
    #
    ThetaInt = integrator(fTheta,rstart,math.log(T0),dr)
    ThetaInt.setParams(params)
    while (params.r < rmax):
        try:
            values = ThetaInt.next(dr)
        except:
            break
        if params.M == 'No Convergence':
            break
        params.theta = math.exp(values[1])
        params.r = values[0]
        rL.append(params.r)
        ML.append(params.M)
        thetaL.append(params.theta)
        TL.append(T(params))
        #Density
        p = params.p0*(TL[-1]/params.theta)**(1./Rcp)
        nL.append((p/(R*TL[-1]))/m)
        #
#
#   Print out terminal results here
    p = params.p0*(TL[-1]/params.theta)**(1./Rcp)
    nbot = (p/(R*TL[-1]))/m
    print 'rmin,pbot,Tbot,nbot,flux = ',rL[-1],p/1.e5,TL[-1],nbot,params.F
#
#   Return results for plotting
#
    c = Curve()
    c.addCurve(rL)
    c.addCurve(ML)

    c1 = Curve()
    c1.addCurve(rL)
    c1.addCurve(TL)
    c1.addCurve(thetaL)

    c2 = Curve()
    c2.addCurve(rL)
    c2.addCurve(nL)
    print "Maximum Mach Number =",max(ML)
    return c,c1,c2


#Set up Newton solver to match lower boundary temperature
def fTbot(ncrit,rcrit):
    return doCalc(ncrit,rcrit,'Backward','Iterate')
solver = newtSolve(fTbot)
solver.setParams(35.)
solver.tolerance = .01
solver.eps = 1.


#gas = AtomicO
#gas = phys.O2
gas = phys.H2
#gas = AtomicH
#gas = Oobleck

Rcp = gas.Rcp
beta = 1.-Rcp
gam = 1./beta
cp = gas.cp
R = gas.R
m = gas.MolecularWeight/(1000.*phys.N_avogadro)
planet = planets.Earth
gs = planet.g
rs = planet.a

#Specify heating rate and depth scale
# (Try Q0 = .001, r1 = 1
# Vary n0 to change shape. If n0 is too small temperature will
# halt at zero before the lower boundary is reached. As n0 is increased
# you see a more pronounced temperature dip. If r1 = 1, you never see
# much of a dip. With r1 = 2, you can get a pretty good dip. Cases
#with a dip tend to have huge density at the base, though
#
Q0 = .001
rh3 = 1.
rh1 = 1.101
rlim = .5*35.

#MachList = [.1 + .01*i for i in range(300)]
MachList = [.01*i for i in range(1,1000)]
rL = [1.,1.1,1.2,1.3,1.4]


def looper(rcStart,ncStart):
    n = ncStart
    for rcrit in [rcStart +5.*i for i in range(60)]:
	solver.setParams(rcrit)
	n = solver(n)
	print '==========Solution:',rcrit,n,"==========================="