#!/usr/bin/env python3
"""Draw a damped harmonic oscillator in matplotlib's xkcd style.

    python3 pendulum_plot.py            # writes pendulum.png and shows it
    python3 pendulum_plot.py -o out.png --no-show

The style call is the novelty; the useful part is everything around it —
annotations, arrows, shaded regions, reference lines and the removal of axis
furniture. Those techniques apply to any matplotlib figure.

Requires NumPy and matplotlib.
"""

import argparse

import numpy as np
import matplotlib.pyplot as plt


def damped_oscillator(n=10_000, span=200.0, frequency=0.5, decay=0.03):
    x = np.linspace(0, span, n)
    return x, np.cos(frequency * x) * np.exp(-decay * x)


def draw(x, y, out_path="pendulum.png", show=True):
    plt.xkcd()

    fig = plt.figure(figsize=(12, 5), dpi=150, facecolor="white")
    axes = plt.subplot(111)

    plt.plot(x, y, color="red", linewidth=2, linestyle="-", zorder=5, alpha=0.5)

    # The exact values are not the point, so the ticks come off entirely.
    axes.set_xticks([])
    axes.set_yticks([])

    excess = 1.15
    axes.set_xlim(excess * x.min(), excess * x.max())
    axes.set_ylim(excess * y.min(), excess * y.max())

    # Mark the settled end of the curve.
    last = -100
    plt.scatter(x[last], y[last], s=50, zorder=12, c="black")
    plt.text(x[last] - 15, y[last] + 0.05, "I'm done here!",
             ha="left", va="bottom")

    plt.ylabel("LEFT      0      RIGHT", y=0.5, fontsize=20)
    plt.text(x[100], y.min() - 0.15, "Time this way->",
             ha="left", va="top", color="black", fontsize=20)

    # Shade above and below the axis. axhspan spans the full width of the
    # axes whatever the limits are; fill_between would stop at the last data
    # point and leave a bare strip on the right, since the x limit is set
    # 15% beyond the data.
    axes.axhspan(0, y.max() * excess, color="0.85", zorder=-1)
    axes.axhspan(y.min() * excess, 0, color=(1.0, 1.0, 0.9), zorder=-1)

    # Bracket the region where the swing is still wild.
    left, right = x[400], x[3000]
    axes.axvline(left, color=".5", ls="--")
    axes.axvline(right, color=".5", ls="--")
    plt.text(x[1400], 0.95, "WOOHOO", ha="left", va="top", fontsize=12)
    plt.text(x[900], 0.8, "craziness zone", ha="left", va="top", fontsize=12)
    plt.annotate("", xytext=(left, 0.85), xy=(right, 0.85),
                 arrowprops=dict(arrowstyle="<->"))

    plt.axhline(color="gray", linewidth=1, zorder=2)
    plt.title("Life of a pendulum", fontsize=24)

    fig.savefig(out_path, bbox_inches="tight")
    print(f"wrote {out_path}")
    if show:
        plt.show()
    plt.close(fig)


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("-o", "--out", default="pendulum.png")
    parser.add_argument("--no-show", action="store_true",
                        help="save without opening a window")
    args = parser.parse_args(argv)

    x, y = damped_oscillator()
    draw(x, y, out_path=args.out, show=not args.no_show)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())