 from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft


def get_fft_values(y_values, T, N, f_s):
    f_values = np.linspace(0.0, 1.0/(2.0*T), N//2)
    fft_values_ = fft(y_values)
    fft_values = 2.0/N * np.abs(fft_values_[0:N//2])
    return f_values, fft_values


t_n = 10
N = 1000
T = t_n / N
f_s = 1/T
#473.04619591757006 0.002594498432061533
#236.0230491258912 0.001201773575026383
#33.003222970993264 0.0008165787374716802
x_value = np.linspace(0,t_n,N)
amplitudes = [5, 7, 9, 12, 17]
frequencies = [6, 4, 2.7, 1., 0.9]

#amplitudes = [0.0008165787374716802, 0.001201773575026383, 0.002594498432061533]
#frequencies = [33.003222970993264, 236.0230491258912, 473.04619591757006]

y_values = [amplitudes[ii]*np.sin(2*np.pi*frequencies[ii]*x_value) for ii in range(0,len(amplitudes))]
composite_y_value = np.sum(y_values, axis=0)

f_values, fft_values = get_fft_values(composite_y_value, T, N, f_s)

colors = ['k', 'b', 'b', 'b', 'b', 'b', 'b', 'b', 'b']

fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111, projection='3d')
#plt.title("")
ax.set_xlabel("\nTime in s", fontsize=16)
ax.set_ylabel("\nFrequency in Hz", fontsize=16)
ax.set_zlabel("\nAmplitude", fontsize=16)

y_values_ = [composite_y_value] + list(reversed(y_values))
frequencies = [0.9, 1., 2.7, 4, 6]

for ii in range(len(frequencies)):
    signal = y_values_[ii]
    color = colors[ii]
    length = signal.shape[0]
    x=np.linspace(0,t_n,N)
    y=np.array([frequencies[ii]]*length)
    z=signal

    if ii == 0:
        linewidth = 4
    else:
        linewidth = 2
    ax.plot(list(x), list(y), zs=list(z), linewidth=linewidth, color=color)

    x=[10]*75
    y=f_values[:75]
    z = fft_values[:75]*3
    ax.plot(list(x), list(y), zs=list(z), linewidth=2, color='red')

    plt.tight_layout()
plt.legend(["Original signal","Frequency with corresponding amplitude","Signal components"])
plt.show()
