import numpy as np
import pandas as pd
import matplotlib.font_manager
import matplotlib.pyplot as plt

matplotlib.use("TkAgg")


def isLocalMax(data: "np.ndarray", idx: "int", win: "int" = 10) -> "bool":
    if idx < win or idx > data.size - win:
        return False
    for i in range(1, win + 1):
        if data[idx] < data[idx - i] or data[idx] < data[idx + i]:
            return False
    return True


def getFreqAndAmple(
        data: "np.ndarray",
        count: "int" = 10,
        hold: "int" = 500,
        win: "int" = 10,
        img: "bool" = False
) -> "np.ndarray":
    time = np.linspace(0, 0.38 * data.size / 128, data.size)
    fft = np.fft.fft(data)
    frq = np.fft.fftfreq(data.size, time[1] - time[0])

    pFrq = frq[:frq.size // 2]
    pFft = np.abs(2.0 / data.size * fft[:fft.size // 2])

    result, idx = np.zeros((count, 2)), 0
    for i in range(win, pFft.size - win):
        if isLocalMax(pFft, i, win) and pFft[i] > hold:
            result[idx][0], result[idx][1] = pFrq[i], pFft[i]
            idx += 1
        if idx >= count:
            break
    print(result)

    if img:
        # 绘制原始信号
        plt.subplot(2, 1, 1)
        plt.plot(time, data)
        plt.title('Raw Frequency')
        plt.xlabel('Frequency (Hz)')
        plt.ylabel('Magnitude')
        # 绘制频率信号
        plt.subplot(2, 1, 2)
        plt.plot(pFrq, pFft)
        plt.title('Positive Frequency')
        plt.xlabel('Frequency (Hz)')
        plt.ylabel('Magnitude')
        plt.show()

    return result


def main():
    file = "data/csv/ac/电动车波形2.csv"
    data = pd.read_csv(file, header=None)
    signal = data[0].to_numpy().flatten()[:128]
    getFreqAndAmple(signal, count=10, hold=180, win=3, img=True)


if __name__ == "__main__":
    main()