import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

# Anscombe's quartet (Anscombe, 1973). Every panel prints its own statistics,
# computed here from the data, to show that they coincide.
x123 = [10, 8, 13, 9, 11, 14, 6, 4, 12, 7, 5]
y1 = [8.04, 6.95, 7.58, 8.81, 8.33, 9.96, 7.24, 4.26, 10.84, 4.82, 5.68]
y2 = [9.14, 8.14, 8.74, 8.77, 9.26, 8.10, 6.13, 3.10, 9.13, 7.26, 4.74]
y3 = [7.46, 6.77, 12.74, 7.11, 7.81, 8.84, 6.08, 5.39, 8.15, 6.42, 5.73]
x4 = [8, 8, 8, 8, 8, 8, 8, 19, 8, 8, 8]
y4 = [6.58, 5.76, 7.71, 8.84, 8.47, 7.04, 5.25, 12.50, 5.56, 7.91, 6.89]

datasets = [('I', x123, y1, 'a line plus noise'),
            ('II', x123, y2, 'a smooth curve'),
            ('III', x123, y3, 'a line plus one outlier'),
            ('IV', x4, y4, 'one point makes r')]

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 4, figsize=(10, 3.1), sharex=True, sharey=True)
for ax, (name, x, y, story) in zip(axes, datasets):
    x, y = np.array(x, float), np.array(y)
    b, a = np.polyfit(x, y, 1)
    r = np.corrcoef(x, y)[0, 1]
    xs = np.linspace(2, 20, 10)
    ax.plot(xs, a + b * xs, '-', color=BLUE, linewidth=1.2)
    ax.plot(x, y, 'o', color=ORANGE, markersize=4.5)
    ax.set_title(f'{name}: {story}', fontsize=9)
    ax.text(0.97, 0.04,
            f'mean x {x.mean():.1f}  mean y {y.mean():.2f}\n'
            f'var x {x.var(ddof=1):.1f}  var y {y.var(ddof=1):.1f}\n'
            f'r {r:.2f}  line {a:.2f} + {b:.2f} x',
            transform=ax.transAxes, va='bottom', ha='right', fontsize=7.5, family='monospace')
    ax.set_xlim(2, 20)
    ax.set_ylim(2, 14.5)
    ax.tick_params(labelsize=7.5)
    ax.set_xlabel('x', fontsize=8)
axes[0].set_ylabel('y', fontsize=8)
fig.suptitle("Anscombe's quartet: the same statistics, four different stories", fontsize=10)

plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()