import sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

CSV_FILE = sys.argv[1] if len(sys.argv) > 1 else "lorenz_member1.csv"

FIT_T_MIN = 2.0
FIT_DELTA_STOP = 1.0

ENSEMBLE_LAMBDAS = [
    0.97226286, 0.85264695, 0.91045011, 0.82656343,
    0.86580801, 0.68093853, 0.61532782, 0.93700653,
]
SPROTT_LAMBDA = 0.906

data = np.genfromtxt(CSV_FILE, delimiter=",", names=True)
t = data["t"]
xA, yA, zA = data["xA"], data["yA"], data["zA"]
xB = data["xB"]
delta = data["delta"]


fig = plt.figure(figsize=(7, 6))
ax = fig.add_subplot(111, projection="3d")
ax.plot(xA, yA, zA, linewidth=0.4, color="#1f4e79")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.set_title("Figure 1: The Lorenz attractor")
fig.tight_layout()
fig.savefig("lorenz_attractor.png", dpi=200)
plt.close(fig)


fig, ax = plt.subplots(figsize=(9, 4))
ax.plot(t, xA, linewidth=0.7, color="#1f4e79", label="Simulation A (unperturbed)")
ax.plot(t, xB, linewidth=0.7, color="#c00000", label="Simulation B (perturbed by 1e-5)")
ax.set_xlabel("time (model units)")
ax.set_ylabel("x")
ax.set_title("Figure 2: Two trajectories from almost-identical starts")
ax.legend(loc="upper right", fontsize=9)
fig.tight_layout()
fig.savefig("lorenz_divergence.png", dpi=200)
plt.close(fig)


mask = np.zeros_like(delta, dtype=bool)
started = False
for i in range(len(t)):
    if not started:
        if t[i] > FIT_T_MIN:
            started = True
        else:
            continue
    if delta[i] >= FIT_DELTA_STOP:
        break
    if delta[i] > 0.0:
        mask[i] = True

tw = t[mask]
lnd = np.log(delta[mask])
slope, intercept = np.polyfit(tw, lnd, 1)

fig, ax = plt.subplots(figsize=(9, 5))

positive = delta > 0
ax.semilogy(t[positive], delta[positive], linewidth=0.7, color="#1f4e79",
            label="separation \u03b4(t)")
fit_line = np.exp(intercept + slope * tw)
ax.semilogy(tw, fit_line, linewidth=2.0, color="#c00000",
            label=f"fitted slope \u03bb = {slope:.3f}")
ax.axvline(tw[0], color="grey", linestyle=":", linewidth=1)
ax.axvline(tw[-1], color="grey", linestyle=":", linewidth=1)
ax.set_xlabel("time (model units)")
ax.set_ylabel("separation \u03b4 (log scale)")
ax.set_title("Figure 3: Exponential growth of the separation")
ax.legend(loc="lower right", fontsize=9)
fig.tight_layout()
fig.savefig("lorenz_separation.png", dpi=200)
plt.close(fig)

lam = np.array(ENSEMBLE_LAMBDAS)
mean, std = lam.mean(), lam.std(ddof=1)

fig, ax = plt.subplots(figsize=(7, 4))
xpos = np.arange(1, len(lam) + 1)
ax.scatter(xpos, lam, color="#1f4e79", zorder=3, label="ensemble members")
ax.axhline(mean, color="#c00000", linewidth=1.5,
           label=f"mean = {mean:.3f} \u00b1 {std:.3f}")
ax.fill_between([0.5, len(lam) + 0.5], mean - std, mean + std,
                color="#c00000", alpha=0.12)
ax.axhline(SPROTT_LAMBDA, color="green", linestyle="--", linewidth=1.5,
           label=f"published value = {SPROTT_LAMBDA}")
ax.set_xlim(0.5, len(lam) + 0.5)
ax.set_xlabel("ensemble member")
ax.set_ylabel("measured \u03bb")
ax.set_title("Figure 4: Lyapunov exponent across the ensemble")
ax.legend(loc="upper right", fontsize=8)
fig.tight_layout()
fig.savefig("lorenz_ensemble.png", dpi=200)
plt.close(fig)

print(f"CSV read: {CSV_FILE} ({len(t)} rows)")
print(f"Figure 3 fit window: t in [{tw[0]:.3f}, {tw[-1]:.3f}], "
      f"member lambda = {slope:.4f}")
print(f"Ensemble: mean lambda = {mean:.4f}, std = {std:.4f} "
      f"(n = {len(lam)})")
print("Saved: lorenz_attractor.png, lorenz_divergence.png, "
      "lorenz_separation.png, lorenz_ensemble.png")