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

import lorenz63 as l63

# BACKGROUND ATTRACTOR

w, t = l63.generate_background_attractor()

# PERTURBED ATTRACTOR

# Time settings
# Forecast base time
t_0 = 120.
# Total integration time
Dt = 1.
# Output timestep
dt = 0.01
# Ensemble size
N = 100

# Initial condition settings
# Como condicion inicial central usamos el estado del sistema de referencia
# al tiempo t_0
w0 = l63.find_system_state(w, t, t_0)
print 'Estado inicial a t = %5.2f: ' % (t_0) + \
    ': (x0, y0, z0) = (%5.2f, %5.2f, %5.2f)' % \
    (l63.unpack_state(w0))
# Estado futuro del sistema despues de un tiempo Dt
w1 = l63.find_system_state(w, t, t_0 + Dt)
print 'Estado final a t = %5.2f: ' % (t_0 + Dt) + \
    ' (x0, y0, z0) = (%5.2f, %5.2f, %5.2f)' % \
    (l63.unpack_state(w1))

# Generacion de las condiciones iniciales para el ensamble
w0_pert = l63.generate_ensemble_initial_conditions(w0, N)

# Pronostico
w_fct, t = l63.integrate_ensemble(w0_pert, Dt, dt)

# El estado futuro real y el pronostico por conjuntos en el espacio fase
plt.figure()
plt.clf()
# Atractor de referencia
plt.plot(w[:, 0], w[:, 2], color='grey')
# Pronostico
plt.plot(w_fct[0, 0, :], w_fct[0, 2, :], marker='o', linestyle='none')
plt.plot(w_fct[-1, 0, :], w_fct[-1, 2, :], marker='o', linestyle='none')
# Estado futuro real
plt.plot(w0[0], w0[2], marker='s', color='black', \
             linestyle='none')
plt.plot(w1[0], w1[2], marker='s', color='black', \
             linestyle='none')
plt.xlabel('x')
plt.ylabel('z')

# La funcion de densidad de probabilidad generado por el ensamble
l63.plot_pdfs(w_fct, w_ana=w1)

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