#This plots the analytic solution to the equation
# dT/dt = 1 - T**4. (Problem set 2)
#
import math
from ClimateUtilities import *

def F(T,T0):
	return math.log( ((T+1.)*(T0-1)/((T-1.)*(T0+1.)))**.25)

def G(T,T0):
	z = (T+1j)/(T-1j)
	z0 = (T0+1j)/(T0-1j)
	return math.log( ((z/z0)**.25j).real)

#Approximate solution for large T
def bigT(T,T0):
    return (1./3.)* (1./T**3 - 1./T0**3)

#Approximate solution for T near fixed point
def fpT(T,T0):
    theta0 = math.atan(1/T0)
    return .25*math.log((2/(T-1))*((T0-1)/(T0+1))) - .5*(math.pi/4 - theta0)

#Approximate exponential solution near fixed point

#Function to return the approximate sol'n curve for a given initial condition
def approx(T0,n):
    tlist1 = []
    tlist2 = []
    Tlist = []
    for T in  [T0 + (1-T0)*i/(n-1) for i in range(n-1)]:
        Tlist.append(T)
        tlist1.append(bigT(T,T0))
        tlist2.append(fpT(T,T0))
    return tlist1,tlist2,Tlist


#Function to return the solution curve for a given initial condition
def solution(T0,n):
    tlist = []
    Tlist = []
    for T in  [T0 + (1-T0)*i/(n-1) for i in range(n-1)]:
        Tlist.append(T)
        tlist.append(F(T,T0)+G(T,T0))
    return tlist,Tlist

#Get some solutions for various initial conditions
t1,T1 = solution(.01,200)
t2,T2 = solution(.1,200)
t3,T3 = solution(.5,200)
t4,T4 = solution(1.5,200)
t5,T5 = solution(2.,200)
#Put them in a curve and plot them
#**ToDo:  Overlay multiple curves on a single plot
c1 = Curve()
c1.addCurve(T1)
c1.addCurve(t1)
c1.switchXY = True
plot(c1)

c2 = Curve()
c2.addCurve(T5)
c2.addCurve(t5)
c2.switchXY = True
plot(c2)

#**ToDo: Compare exponential solution
#Now do the approximate solutions and compare, in one of the cases
tbig,tfp,T = approx(2.,200)
c3 = Curve()
c3.addCurve(T5)
c3.addCurve(t5,'t','Exact')
c3.addCurve(tbig,'tb','Large T approx')
c3.addCurve(tfp,'tfp','Linearized approx')
c3.switchXY = True
plot(c3)