#!/usr/bin/env python
import matplotlib.pyplot as plt

import lorenz63 as l63

# BACKGROUND ATTRACTOR

w, t = l63.generate_background_attractor()

# PERTURBED ATTRACTOR

# Time settings
# Total integration time
t1 = 2.
# Output timestep
dt = 0.01

# Initial condition settings
# Central initial condition
# [-10.06, -12.49, 25.75) => small spread; 
# [1.21, 1.02, 18.54] => larger spread
x0 = -10.06
y0 = -12.49
z0 = 25.75

# Ensemble size
N = 200

w0_pert = l63.generate_ensemble_initial_conditions([x0, y0, z0], N)

w_pert, t = l63.integrate_ensemble(w0_pert, t1, dt)

plt.figure()
plt.clf()

plt.plot(w[:, 0], w[:, 2], color='grey')

plt.plot(w_pert[0, 0, :], w_pert[0, 2, :], marker='o', linestyle='none')
plt.plot(w_pert[-1, 0, :], w_pert[-1, 2, :], marker='o', linestyle='none')

plt.xlabel('x')
plt.ylabel('z')

plt.show()
plt.close('all')