# -*- coding: utf-8 -*-
"""
Created on Mon Feb 04 18:07:47 2019

@author: Alex
"""

from PyDSTool import *
from copy import copy
from matplotlib import pyplot as plt


pars = {
        'CPsea': 250., 
        'CNsea': 37.5,
        'CPland': 1.,
                
	#Taille initiale des réservoirs
        'P0': 0.0031,
		'A0': 3.193,
		'T0': 288.,

	#Constantes Mills

		'k1': 9.,
		'k2': 15.*10**(-3),
		'k3': 6.*10**(-3),
		'k4': 0.86,
		'k5': 6.65,
		'k6': 13.35,
		'k7': 7.75,
		'k8': 57.*10**(-3),
		'k9': 1.25,
		'k10': 6.65,


		'eps': 0.773,
		'S': 1298.,
		'sigma': 5.67*10**(-8),
		'mu': 0.1815,
		'ca': 0.2507,
		'cb': 3.9216,
		'cc': 23.8113,
		'cd': 83.4113,
		'ce': 131.6138,
		'g': 0.78,
		'k_ice': -2,
		'rho': 5.,
		'W_max': 1.1,
		'T_crit': 283.,
		'T_escape': 263.,
		'c': 1158.,

		'ocdeg': 1.25,
		'ccdeg': 6.65,
		'gammma': 0.2,
        }

auxvars = ['var1','var2']


var1 = 'alpha2(T)'
var2 = 'phi(Ab,T)'

icdict = {'P': 1., 'Ab': 1., 'T': 0.}
t_end = 500


P_dot = '(phosw(Ab,T) - fepb(P) - capb(P) - mopb(P))/P0'

A_dot = '(oxidw(T) + ocdeg + ccdeg - mocb(P) - silw(Ab,T))/A0'

T_dot = '((Ri(T)-R0(T))/c + mu*(F_CO2(Ab,T)/F_CO2(1.,T) - 1.))*1000'


DSargs = args(name='COPSE_cycles')
DSargs.pars = pars


DSargs.varspecs = {'P': P_dot, 'Ab': A_dot, 'T': T_dot, 'var1': var1, 'var2': var2}


fnspecs  = {
                    # Fonctions Phosphore
                    'fepb': (['P'], '(k3/k4)*(1.-anox(P))*(mocb(P)/mocb(1.))'),
                    'capb': (['P'], 'k2 * mocb(P)/mocb(1.)'),
		  			'mopb': (['P'], 'mocb(P)/CPsea'),
                    
                    # Fonctions Oxygène
                    'mocb': (['P'], 'k1*(P)**2'),
                    'oxidw': (['T'], 'k7*W_glacial(T)'),
                    
                    # autres fonctions nécessaires
                    'phosw': (['Ab','T'], 'k8*((5./12.) + (1./6.)*silw(Ab,T)/silw(1.,T0) + (5./12.))'),
					'anox': (['P'], '1.-k4*(1./P)'),
                    
                    # Fonctions liées à la température
                    'silw': (['Ab','T'], 'k5*W(Ab,T)'),
					'W': (['Ab','T'], '(W_max-W_K(Ab,T))/(1.+exp(-100.*(W_K(Ab,T)-W_max))) + W_K(Ab,T)'),
					'W_K': (['Ab','T'], 'W_glacial(T)*rho*exp(7537.69*(T+273.-T0)/((T+273)*T0))*(exp(0.03*(T+273.-T0)))**0.65 * (phi(Ab,T)/phi(1.,T0))**0.5'),
					'Ri': (['T'], 'S*(1-alpha2(T))'),
					'R0': (['T'], '4*sigma*((T+273.)**4) *(1-eps/2)'),
					'W_glacial': (['T'], '0.501 - 0.5*atan(10.*(alpha2(T)-0.5))*(2/3.14)'),
					#'W_glacial': (['T'], 'if(T>0, 1., 0.01)'),
					'F_CO2': (['Ab','T'], 'ca*lambdb(Ab,T)**4 + cb*lambdb(Ab,T)**3 + cc*lambdb(Ab,T)**2 + cd*lambdb(Ab,T) + ce'),
					'phi': (['Ab','T'], 'g*(Ab*A0)/((Ab*A0)+h(T))'),
					'lambdb': (['Ab','T'], '(log(phi(Ab,T)))/log(10)'),
					#'h': (['T'], 'if(T<0, 5.*10**2, 1.*10**2)'),
					'h': (['T'], '3.*10**2 + 2.*10**2 * atan(50.*(alpha2(T)-0.5))*(2/3.14)'),
					'alpha2': (['T'], '0.5 - atan(gammma*(T-2))/8.'),
                    }


DSargs.fnspecs  = fnspecs
DSargs.auxvars = auxvars


DSargs.ics = icdict
DSargs.tdata = [0,t_end]
DSargs.algparams = {'init_step': 0.01}

print('Computing graphes...')
start = clock()

simpleDS = Generator.Radau_ODEsystem(DSargs)
traj = simpleDS.compute('COPSE_cycles')
pts  = traj.sample()


tog1 = 2

if tog1 == 1:
	plt.figure(1)
	plt.plot(pts['t']/2, pts['P'], color = 'r')
	plt.title('Evolution temporelle du Phosphore')
	plt.xlabel('Temps (Myr)')
	plt.ylabel('Phosphore')


	fig2, ax1=plt.subplots()
	ax2=ax1.twinx()
	ax1.plot(pts['t']/2, pts['T'], color = 'c')
	plt.title('Evolution temporelle de la temperature et l albedo')
	ax1.set_xlabel('Temps (Myr)')
	ax1.set_ylabel('Temperature (C)', color = 'c')
	ax2.plot(pts['t']/2, pts['var1'], color = 'g')
	ax2.set_ylabel('albedo', color = 'g')
	ax2.axis([0,250,0,1])


	plt.figure(4)
	plt.plot(pts['T'], pts['Ab'], color = 'g')
	plt.title('Portrait de phase du Carbone et de la temperature')
	plt.xlabel('Temperature (C)')
	plt.ylabel('Carbone')


	fig4, ax1=plt.subplots()
	ax2=ax1.twinx()
	ax1.plot(pts['t']/2, pts['var2'], color = 'k')
	plt.title('Evolution temporelle du reservoir de Carbone et du CO2 atmospherique')
	ax1.set_xlabel('Temps (Myr)')
	ax1.set_ylabel('CO2 (atm)', color='k')
	ax1.axis([0,250,0,0.4])
	ax2.plot(pts['t']/2, pts['Ab'], color = 'g')
	ax2.set_ylabel('Carbone', color='g')


print('done in %.3f seconds!' % (clock()-start))

if tog1 == 2:
	PyCont = ContClass(simpleDS)

	PCargs = args(name='EQ1', type='EP-C')
	PCargs.freepars = ['rho']
	PCargs.StepSize = 1e0
	PCargs.MaxStepSize = 1e0
	#PCargs.MinStepSize = 1e-2
	PCargs.MaxNumPoints = 50
	PCargs.LocBifPoints = 'all'
	PCargs.verbosity = 2
	PCargs.SaveJacobian = True
	PCargs.SaveEigen = True
	PyCont.newCurve(PCargs)

	print('Computing curve...')
	start = clock()
	PyCont['EQ1'].forward()
	#PyCont['EQ1'].backward()
	print('done in %.3f seconds!' % (clock()-start))

	PCargs = args(name='LCQ1', type='LC-C')
	PCargs.freepars = ['rho']
	PCargs.initpoint = 'EQ1:H1'
	PCargs.StepSize = 1e-1
	PCargs.MaxStepSize = 1e-1
	PCargs.MaxNumPoints = 100
	PCargs.NumSPOut = 100
	#PCargs.LocBifPoints = 'all'
	PCargs.SolutionMeasures = 'min'
	PCargs.verbosity = 2
	PCargs.SaveJacobian = True
	PCargs.SaveEigen = True
	PyCont.newCurve(PCargs)

	print('Computing curve...')
	start = clock()
	PyCont['LCQ1'].backward()
	print('done in %.3f seconds!' % (clock()-start))

	plt.figure(1)

	PyCont['EQ1'].display(('rho','T'), stability=True)
	PyCont['LCQ1'].display(('rho','T'), stability=False)

plt.show()