import numpy as np
import matplotlib.pyplot as plt

plt.rc('text', usetex=True)

plt.figure(figsize=(4,3))

s = 5

x = np.arange(0, s, 1.0/256)
y = x*1
y[-1] = 0

steps = np.zeros_like(x)

for k in range(1,s+1):
    steps[x>k-1] = k-.5
    steps[x==k] = 0


plt.title('converting sum to integral where $s=%s$' % s)

plt.plot(x, y, label=r'$x$')
plt.plot(x, steps, label=r'$k-\frac12$')
plt.xlabel('$x$')

plt.legend(loc='best')

plt.ylim(0,s)

plt.tight_layout()

# plt.savefig('../web/figs/sum-gives-square', dpi=600, format='png')
plt.savefig('sum-gives-square.svg')
plt.savefig('sum-gives-square.pdf')

plt.figure(figsize=(4,3))
plt.title('solving sum geometrically where $s=%s$' % s)

steps[-1] = 0
plt.plot([s/2, s/2], [0, s])
plt.plot(x, steps, label=r'$k-\frac12$')
plt.plot(s-x[x>=s/2], s-steps[x>=s/2], 'g-')
plt.plot(x[x>=s/2], steps[x>=s/2], 'g--')
plt.xlabel('$x$')

plt.legend(loc='best')

plt.ylim(0,s)

plt.tight_layout()

# plt.savefig('../web/figs/sum-geometrically', dpi=600, format='png')
plt.savefig('sum-geometrically.svg')
plt.savefig('sum-geometrically.pdf')