# lorenzSDIC.py
"""
  Plot of two Lorenz 3D curves, with separate plots of their
  x values over time
"""

import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from frange import *

DT = 0.01
N = 10000


def lorenz(xyz1, s=10, p=28, b=8/3):
  x, y, z = xyz1
  x_dot = s*(y - x)
  y_dot = p*x - y - x*z
  z_dot = x*y - b*z
  return (x_dot, y_dot, z_dot)


xyz1 = [(0,0,0) for x in range(N+1)]
xyz1[0] = (0, 1.0, 1.05)

delta=0.001
xyz2 = [(0,0,0) for x in range(N+1)]
xyz2[0] = (delta, 1.0, 1.05)

for i in range(N):
  xyz1[i+1] = tuple(map(lambda x,y: x+y*DT, 
                                xyz1[i], lorenz(xyz1[i])))

  xyz2[i+1] = tuple(map(lambda x,y: x+y*DT, 
                                xyz2[i], lorenz(xyz2[i])))

fig = plt.figure(figsize=(12, 5))
gs = GridSpec(2, 2)  # no. rows, cols
gs.update(wspace=0.3, hspace = 0.3)

ax1 = fig.add_subplot(gs[:, 0:1], projection='3d')  
              # occupies first column
ax2 = plt.subplot(gs[0, 1:])  # rest of first row
ax3 = plt.subplot(gs[1, 1:])  # rest of 2nd row

# 3D plot of two curves
xs1, ys1, zs1 = zip(*xyz1)
ax1.plot(xs1, ys1, zs1, lw=0.5, color="green", alpha=0.5,
                        label="x10="+str(xyz1[0]))
xs2, ys2, zs2 = zip(*xyz2)
ax1.plot(xs2, ys2, zs2, lw=0.5, color="red", alpha=0.5,
                        label="x20="+str(xyz2[0]))
ax1.set_xlabel("X")
ax1.set_ylabel("Y")
ax1.set_zlabel("Z")
ax1.set_title("Lorenz Attractor")

ts = [i for i in range(N+1)]

#---- First time plot
ax2.plot(ts, xs1, color="green")
ax2.axhline(0, color='k', lw=.5, alpha=.5)
ax2.set_xticks([])
ax2.set_ylabel('x')
ax2.set_title(str(xyz1[0])+" vs. time")

#---- second time plot
ax3.plot(ts, xs2, color="red")
ax3.axhline(0, color='k', lw=.5, alpha=.5)
ax3.set_xticks([])
ax3.set_ylabel('x')
ax3.set_title(str(xyz2[0])+" vs. time")

plt.show()