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

def p_func(x, a, om, ph):
    """
    a sind die verschiedenen Amplituden
    om sind die verschiednenen Freuquenzen
    ph sind die Phasen
    """
    y = np.zeros(len(x))
    for ampl, omega, phi in zip(a, om, ph):
        phase = phi*np.ones(len(x))
        y+= ampl*np.cos(omega*x + phase)
    return y

def ind_func(x, a, om, phi):
    phase = phi*np.ones(len(x))
    return a*np.cos(om*x + phase)


n = 6
nax = np.arange(6)
a = np.ones(n)
om = np.array([1., 2., 3., 4., 5.])
ph = np.array([0., 0., 0., 0., 0.])
x = np.linspace(0, n*np.pi, 1000)

fig, axs = plt.subplots(6, sharex = True)
for ax, i in zip(axs, nax):
    print(i)
    if i < n-1:
        ax.plot(x, ind_func(x, a[i], om[i], ph[i]), 'k-')
        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])
    if i == n-1:
        ax.set_ylim(-2, n)
        ax.plot(x, p_func(x, a, om, ph), 'r-')
        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)
plt.subplots_adjust(hspace=0)
plt.savefig('fourier_1.eps', bbox_inches='tight')
plt.show()