#
#Use this script as the basis of the basic linearized EBM
#for the Meridional Heat Transport chapter (Chapter 9)
#This is a steady state EBM, without seasonal cycle or ice
#growth
#
#
#Data on section of text which this script is associated with
Chapter = '9.**'
Figure = '**'
#


from ClimateUtilities import *
from math import *

# Define functions
#
#Annual mean distribution of solar flux (multiply by Lsun
def sflux(sinlat):
    x = sinlat*sinlat
    return .30584 + x*(-.16542 + x*(.14449 +x*(-.37829 + .22003*x)))

#Slope function defining the differential equation we are solving
def d(y,v,y_ice):
    if y > y_ice:
      alb = alpha_ice
    else:
      alb = alpha_o
    dv = Numeric.zeros(4,'float')
    dv[0] = (A+B*(v[1]-T1) -Lsun*(1.-alb)*sflux(y))/D
    dv[1] = v[0]/(1.-y*y) #v[1] is temperature
# Homogeneous solution
    dv[2] =  B*v[3]/D
    dv[3] = v[2]/(1.-y*y)
    return dv

def doCalc(y_ice):
    #Arrays
    vstart = Numeric.zeros(4,'float')

    #---Setup finished-----
    #-------------Beginning of calculation of temperature profile---

    #Set initial conditions
    vstart[0] = 0.
    vstart[1] = 300.
    vstart[2] = 0.
    vstart[3] = 1. # Homogeneous solution

    #Integration loop
    #Note that f_i.x is the current value of the independent
    #  variable, i.e. the current value of y.  Sorry if this
    #  is confusing. I'm thinking of ways to redesign the
    #  integrator object to reduce this kind of confusion
    #
    f_i = integrator(d,0.,vstart,dy0)
    f_i.setParams(y_ice)
    while f_i.x < .99:
        f_i.dx = dy0*(1.-f_i.x*f_i.x + 1.e-5) #Adjust increment
        ans = f_i.next() #Note:
                         #  ans[0] is y
                         #  ans[1] is the current value of v
                         #      (remember v is a 4-element vector)

    v = ans[1] 
    c1 = -v[0]/v[2]

    #Re-do with correct superposition
    vstart[0] = 0.
    vstart[1] = 300.
    vstart[2] = 0.
    vstart[3] = 1.#Homogeneous solution

    #Lists in which to save results
    latList = []
    TList = []
    #Integration loop
    f_i = integrator(d,0.,vstart,dy0)
    f_i.setParams(y_ice)
    while f_i.x < .99:
        f_i.dx = dy0*(1.-f_i.x*f_i.x + 1.e-5)
        ans = f_i.next()
        v = ans[1]
        T = v[1] + c1*v[3]
        latList.append(f_i.x)
        TList.append(T)

    #Put in Curve for output and plotting
    c = Curve()
    c.addCurve(latList,'sinlat','Sine Latitude')
    c.addCurve(TList,'T','Surface Temperature')
    #Find the temperature at the ice margin
    nice = numpy.searchsorted(latList,y_ice)
    #Do interpolation to find T(y_ice)
    if nice > 0:
        y1,T1 = latList[nice-1],TList[nice-1]
        y2,T2 = latList[nice],TList[nice]
        Ti = T1 + (y_ice-y1)*(T2-T1)/(y2-y1)
    else:
        Ti = TList[0]        
    return Ti,c

def IceMarginT():
    yL = [i/100. for i in range(100)]
    Ti = []
    for y in yL:
        T,c = doCalc(y)
        Ti.append(T)
    c = Curve()
    c.addCurve(yL,'y_IceMargin')
    c.addCurve(Ti,'T_IceMargin')
    c.addCurve([Tfreeze for x in Ti],'T_freeze') #Freeze temperature
    return c

#------------------Set Constants --------------------------------
# OLR constants for linearization around 273K
#               OLR = A + B*(T-T1)
T1 = 260.
A = 202.263
B= 2.32

#Solar constant
Lsun=1370.

#Other Constants
Tfreeze = 271. #Freezing point of sea water (just used for graphics)

#Diffusivity,ice-margin,solar constant
#If the solar constant and the OLR constants
#are given dimensionally, than D has the
#units (W/m**2)/K.  E.g with D=1,
#a 1K pole-eq temperature difference would
#cause a flux sufficient to balance a
#1W/m**2 equatorial or polar radiative imbalance
#
D = 1.
y_ice=.75


#Ice and ocean albedo (Edit these to set the value you want)
alpha_ice = .6 # Do uniform albedo case first
alpha_o = .15 #ERBE clear sky value is .16, including ice and snow
              #Using the clear-sky value should make things too warm
              #That could justify using a somewhat higher albedo

#Stepping constants
dy0 = .0025

#----------------------------------------------------
#The following is an example of how to run the model. 
#
Ti1,c1 = doCalc(.5) #Get ice margin T and T(y) for y_ice = .5
plot(c1) #Plot it
D = .25 # Change diffusivity and do it again
Ti2,c2  = doCalc(.5)
#Now stuff both results into a Curve object and
#plot. (Eventually, I'll provide a more convenient way
#to composite plots)
cAll = Curve() #A new Curve object to plot both
cAll.addCurve(c1['sinlat'],'sinlat')
cAll.addCurve(c1['T'],'D=1')
cAll.addCurve(c2['T'],'D =.25')
plot(cAll)

#Make a graph of ice margin temperature vs. ice margin position
c3 = IceMarginT()
plot(c3)