from math import pi
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def tria(n, om, t):
    return 2./np.pi*np.arctan(1./np.tan(t+np.pi/2.))
    #return -2./np.pi * np.arcsin(np.sin(2.*t))
#    return approximation(500, om, t)

def approximation(n, om, t):
    y = np.zeros(len(t))
    for i in 1+np.arange(n):
        print(i)
        #y+=8./np.pi**2 * (-1)**i/i**2 * np.sin(om*i*t)
        y+=(-1)**i * np.sin(i*om*t)/i
    return 2.*y/np.pi

n = 6
nax = np.arange(6)
print(nax)
a = np.ones(n)
om = 2.
x = np.linspace(0, n*np.pi, 1000)
y = tria(n, om, x)

fig, axs = plt.subplots(6, sharex = True)
for ax, i in zip(axs, nax+1):
    print(i)
    ax.plot(x, y, 'r-')
    ax.plot(x, approximation(i, om, x), 'k-', label=r'n={}'.format(i))
    ax.set_ylim(-1.3, 1.3)
    ax.set_xticks([0., .5*np.pi, pi, 1.5*np.pi, 2.*pi, 2.5*np.pi, 3.*np.pi, 3.5*np.pi, 4.*np.pi, 4.5*np.pi, 5.*np.pi, 5.5*np.pi, 6.*np.pi])
    ax.set_xticklabels([ '0', r'$\frac{\pi}{2}$', r'$\pi$', r'$\frac{3}{2}\pi$', r'$2\pi$', r'$\frac{5}{2}\pi$', r'$3\pi$', r'$\frac{7}{2}\pi$', r'$4\pi$', r'$\frac{9}{2}\pi$', r'$5\pi$', r'$\frac{11}{2}\pi$', r'$6\pi$'], fontsize=14)
    ax.text(5.6*np.pi, -0.9, 'n={}'.format(i), fontsize=14)
    #ax.legend(loc='lower right')
plt.subplots_adjust(hspace=0)
plt.savefig('triangle.eps', bbox_inches='tight')
plt.show()