from numpy import array, random, kaiser, zeros, cos, angle, absolute
from numpy.core.umath import pi
from numpy.fft import fft
import numpy
from matplotlib.pyplot import figure, title, plot, xlim, show
import matplotlib.pyplot as plt
#
#       генерация тестового сигнала (a test signal generation)
#

# параметры сигнала (signal parameters)
n = 1024
harmonics = array([50])
amplitudes = array([1])
phases = array([pi/2])

# генерация (generation)
signal = zeros(n)
noise = 0*random.normal(0, 0.5, n)   # мощный шум (powerful noise)
for i in range(n):
    for j in range(len(harmonics)):
        signal[i] = signal[i] + amplitudes[j] * cos(2 * pi * harmonics[j] * i / n + phases[j])
signal = signal+noise
spectrum = fft(signal)      # для тестовых целей (for test purpose)

#
#       основная часть, анализ сигнала (main part, analise of signal)
#

#       параметры алгоритма (algorithm parameters)
# параметры окна (window parameters)
kaiser_beta = 14            # окно кайзера бета (kaiser window beta)
exp_ratio = 2048            # степень расширения (expansion ratio)
delta = 5                   # дельта окна для обнаружения пика (delta of window for peak detection)
threshold_level = 0.01      # гармонический порог (harmonic threshold)
width_window = 10           # ширина окна (width of window)
# параметры нахождения пика (peak finding parameters)
peak_group_distance = 5     # найти кандидата на пик (to find candidate to peak)
peak_area_distance = 100    # найти максимум пика (to find maximum of peak)

### Оконное преобразование
# calculate window (окно расчета)
window = kaiser(n+1, kaiser_beta)   # создать симметричное окно (create symmetric window)
window_factor = pi*sum(window)/n    # коэффициент окна (window factor)
# применить окно, чтобы уменьшить количество пиков (apply window to reduce a number of peak)
signal2 = signal * window[0:n]
spectrum2 = fft(signal2)

figure(1)
title("signal+noise")
plot(signal)

figure(2)
title("spectrum")
plot(20*numpy.log10(numpy.abs(spectrum)))

figure(3)
#title("signal*window")
plt.xlabel('Время [мс]')
plt.ylabel('Амплитуда сигнала')
plt.grid()
plot(signal2)

figure(4)
#title("spectrum with window")
plt.xlabel('Частота в герцах [Гц]')
plt.ylabel('Амплитуда [дБ]')
plt.xlim(-3, 100)
plt.ylim(-100, 50)
plt.grid()
plot(20*numpy.log10(numpy.abs(spectrum2)))

show()