#!/usr/bin/env python3

# integrates the Lorenz model and calculates the correlation dimension 

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import ode
plt.rcParams.update({'font.size':14})
import nolds

# DEFINE FUNCTIONS
# the model itself
def model(t,x,params):
    Pr,b,r = params
    xdot = np.zeros(3)
    xdot[0] = Pr*x[1] - Pr*x[0]
    xdot[1] = -x[0]*x[2] + r*x[0] - x[1]
    xdot[2] = x[0]*x[1]-b*x[2]

    return xdot

# DEFINE PARAMETERS - DON'T CHANGE
Pr = 10  
b = 8/3 
r = 28 

# PARAMETERS YOU CAN CHANGE FOR CORRELATION DIMENSION CALCULATIONS
# set up start and end times of the integration
tstart = 0
tend = 200

# timestep 
dt = 0.2

# embedding dimension
embedding_dimension = 2

# initial condition 
x0 = [0.01, 0.01, 0.01]

# INTEGRATE AND PLOT
t = np.linspace(tstart,tend,int((tend-tstart)/dt))
x = np.zeros((len(t),3))

x[0,:] = x0 

model_instance = ode(model).set_integrator("dop853")
model_instance.set_initial_value(x[0,:]).set_f_params([Pr,b,r])

for i in range(1,len(t)):
    model_instance.integrate(model_instance.t+dt)
    x[i,:] = model_instance.y


plt.figure()
plt.plot(t,x[:,0])
plt.xlabel('Time')
plt.ylabel('x')

plt.figure()

# function returns a fit and the raw data, we save both
correlation_dimension,data = nolds.corr_dim(x[:,0],embedding_dimension,debug_data=True)


plt.scatter(data[0],data[1],label='Data')
fit = np.poly1d(data[2])
plt.plot(data[0], fit(data[0]),color='red',label='Fit')
plt.xlabel('log r')
plt.ylabel('log C(r)')

print('Correlation dimension = '+str(correlation_dimension))

plt.show()