from __future__ import division # In case someone wants to use Python 2
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns # Just for the colors

mpl.rcParams['axes.labelsize'] = 20
mpl.rcParams['axes.titlesize'] = 22
mpl.rcParams['legend.fontsize'] = 18
mpl.rcParams['xtick.labelsize'] = 14
mpl.rcParams['ytick.labelsize'] = 14
mpl.rcParams['lines.linewidth'] = 3

def dirichlet(n, x):
    return np.where(x!=0, np.sin((n + .5)*x)/np.sin(.5*x), 2*n + 1)

points = 2**10 + 1
kernels = 4
x_min = -np.pi
x_max = np.pi

x_vals = np.linspace(x_min, x_max, num=points)
n_vals = np.arange(1, kernels + 1)
D = dirichlet(np.resize(n_vals, (points, kernels)).T,
              np.resize(x_vals, (kernels, points)))

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
for kernel, n in zip(D, n_vals):
    ax.plot(x_vals, kernel,
            label=r'$\frac{\sin((' + str(n) + r'+1/2)x)}{\sin(x/2)}$')

ax.xaxis.set_ticks(np.pi*np.arange(-4, 5)/4)
ax.xaxis.set_ticklabels([r'$-\pi$', r'$-\frac{3\pi}{4}$', r'$-\frac{\pi}{2}$',
                         r'$-\frac{\pi}{4}$', r'$0$', r'$\frac{\pi}{4}$',
                         r'$\frac{\pi}{2}$', r'$\frac{3\pi}{4}$', r'$\pi$'])
ax.legend()
ax.set_xlim([x_min, x_max])
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$D_n(x)$')
ax.xaxis.grid(color='gray')
ax.yaxis.grid(color='gray')
plt.tight_layout()
plt.savefig('dirichlet.svg')