from math import *
from solar import *
from ClimateUtilities import *
import phys

#The following module gives the polynomial OLR fits.
#It is in the Chapter4Scripts directory, and to be
#used here you either need to add that directory
#to the PYTHONPATH, or put a link to the script in
#your modules directory, or just copy the script over.
import OLRPoly

#-------Define the OLR function------------------------------
def OLRBare(T):
    return phys.sigma*T**4


def OLR1(T):
    return OLRPoly.OLRT(CO2,T)

OLR = OLR1

def f(t,T,params):
    #
    albedo = params.albedo
    hflux = params.hflux
    H = params.H
    lat = params.lat
    #
    S = (1.-albedo)*L*solar(lat,2.*pi*t/year,obliquity)
    return (S - OLR(T) + hflux)/(rho*H*Cp)

#Function that integrates the ODE and computes seasonal cycle
def doCalc(lat,hflux):
    dt = year/100.
    fi = integrator(f,0.,Tstart,dt) # Initial temperature is Tstart
    #
    params = Dummy()
    fi.setParams(params)
    params.H = Ho
    params.albedo = albedo_o
    params.lat = lat
    params.hflux = hflux
    #
    time = []
    T = []
    #First do a 50 year integration to reach equilibrium.
    #You may need to increase this if the mixed layer is
    #very deep, or when trying to find the transition to
    #ice-free conditions, in which case a difference of
    #a half degree or so can throw you into a different regime.
    nyears = 50
    n = int(nyears*year/dt)

    for i in range(n):
        t,TT = fi.next()
        if doIceAlbedo:
            if TT < Tfreeze:
                params.albedo = albedo_i
                params.H = Hi
            else:
                params.albedo = albedo_o
                params.H = Ho

    #Now continue for 1 year and save the results
    for i in range(int(year/dt)):
        result = fi.next()
        time.append(result[0]/year - 50.)
        T.append(result[1])
        TT = result[1]
        if doIceAlbedo:
            if TT < Tfreeze:
                params.albedo = albedo_i
                params.H = Hi
            else:
                params.albedo = albedo_o
                params.H = Ho
    return time,T

#Set up basic parameters
L = 1370.
year = 365.*24.*3600.
degrad = pi/180.
obliquity = 23.*degrad
rho = 1000.
Cp = 4218. #Specific heat for liquid water
CO2 = 300.
Tfreeze = 271.
albedo_o = .2
albedo_i = .6
Ho = 50. #Mixed layer depth for ocean. Usually set to 50 for ocean
Hi = 1.  #Mixed later depth for ice
#
Tstart = 300. #Starting temperature; only matters when we have ice feedback

#Do the calculation for a list of latitudes and
#corresponding dynamical heat fluxes. The dynamical heat
#flux hflux is the amount added or taken away from a column
#by meridional fluid dynamic heat transport
#Standard set of values used for waterworld calculation
#in text
lats = [90.,60.,45.,30.,0.]
hfluxes = [80.,40.,0.,-20.,-35]
doIceAlbedo = False #Flag for whether or not to do ice-albedo feedback
#Loop over latitudes, with varying hfluxes
c = Curve()
#There's really gotta be a better way to write
#this kind of loop! (Looping over synchronous lists)
#You could also make the albedo latitude-dependent here.
n = len(lats) #Number of values to loop over
for i in range(n):
    lat = lats[i]
    hflux = hfluxes[i]
    tag = 'T%f'%lat
    time,T = doCalc(lat*degrad,hflux)
    if i == 0:
        c.addCurve(time,'t','Time (years)')   
    c.addCurve(T,tag)    
plot(c)


#Loop over CO2 with constant latitude
#Values used to illustrate inhibition of polar ice formation
doIceAlbedo = True #Flag for whether or not to do ice-albedo feedback
albedo_o = .33 #Allows for low cloud cover in summer
Ho = 50.
CO2vals = [(2**i)*150. for i in range(6)]
lat = 80.
hflux = 110.
c1 = Curve()
n = len(CO2vals) #Number of values to loop over
for i in range(n):
    CO2 = CO2vals[i] #Uncomment to loop over CO2
    tag = '%fppmv'%CO2
    time,T = doCalc(lat*degrad,hflux)
    if i == 0:
        c1.addCurve(time,'t','Time (years)')   
    c1.addCurve(T,tag)    
plot(c1)

#Remark: Note that with the oversimplified ice-albedo feedback
#used here, the ice actually makes the summer warmer, since the
#ice warms up much faster in response to the increasing sunlight.
#That's an unrealistic artifact of not accounting for the energy
#needed to melt the ice. Without the delay, less of the year is
#subfreezing when you allow ice than when ice is not allowed.
#Is there a simple fix for this?