import numpy as np
import matplotlib.pyplot as plt

class EngineCurrentAnalyse1:
    def __init__(self, fs=50000, base_freq=50, sample_len_in_s=1):
        self.period_samples = fs//base_freq
        self.fs = fs
        self.base_freq = 50
        self.sample_len_in_s = sample_len_in_s
    
    def breakOneBar(self, t, final_signal):
        # 处理后的信号
        processed_signal = np.copy(final_signal)

        for i in range(0, len(t), self.period_samples):
            # 每个周期的前 1/6 采样点
            start_idx = i
            end_idx = i + int(self.period_samples / 6)
            if end_idx > len(t):
                break
            # 生成 cos 函数
            cos_coeff = np.cos(2 * np.pi * self.base_freq*0.97 * t[start_idx:end_idx])
            # 乘以 cos 系数
            processed_signal[start_idx:end_idx] *= cos_coeff
        return processed_signal

    def GenCurrentData(self):
        # 定义参数
        fs = self.fs  # 采样频率 (Hz)
        t = np.linspace(0, self.sample_len_in_s, fs, endpoint=False)  # 时间向量，时长 1 秒

        # 生成 50Hz 信号
        f_signal = self.base_freq  # 信号频率 (Hz)
        signal_50Hz = np.sin(2 * np.pi * f_signal * t)

        # 生成白噪声
        noise = np.random.normal(0, 0.1, t.shape)

        # 生成 2000Hz 变频器白噪声
        f_inverter = 2000  # 变频器频率 (Hz)
        inverter_signal = np.sin(2 * np.pi * f_inverter * t)
        inverter_noise = np.random.normal(0, 0.05, t.shape)
        inverter_noisy_signal = inverter_signal + inverter_noise

        # 叠加信号
        final_signal = signal_50Hz + noise + inverter_noisy_signal
        return t, final_signal, fs, 

    def draw_signal(self, final_signal, t, fs):
        # 进行傅里叶变换
        fourier_transform = np.fft.fft(final_signal)
        # 获取频率轴
        freqs = np.fft.fftfreq(len(final_signal), 1/fs)

        # 取单边频谱并计算幅度（绝对值）
        n = len(final_signal)
        single_sided_fft = 2.0/n * np.abs(fourier_transform[:n//2])
        single_sided_freqs = freqs[:n//2]

        # 将幅度转换为 dB 模式
        single_sided_fft_db = 20 * np.log10(single_sided_fft)

        # 绘制时域信号
        plt.figure(figsize=(12, 10))
        plt.subplot(3, 1, 1)
        plt.plot(t, final_signal, label='Final Signal')
        plt.xlabel('Time (s)')
        plt.ylabel('Amplitude')
        plt.title('Three - phase Motor Line Current Single - phase Waveform with Noise (Time Domain)')
        plt.legend()
        plt.grid(True)

        # 绘制频域信号（dB 模式）
        plt.subplot(3, 1, 2)
        plt.plot(single_sided_freqs, single_sided_fft_db, label='Frequency Domain (dB)')
        plt.xlabel('Frequency (Hz)')
        plt.ylabel('Magnitude (dB)')
        plt.title('Three - phase Motor Line Current Single - phase Waveform with Noise (Frequency Domain in dB)')
        plt.legend()
        plt.grid(True)

        # 绘制频域信号（dB 模式）
        # 绘制频域信号（dB 模式），只占据窗口宽度的 1/3
        # 计算合适的位置和大小参数
        span=0.25
        left = (1-span)/2  # 子图左边界位置
        bottom = 0.05  # 子图下边界位置
        width = span  # 子图宽度，占据窗口宽度的 1/3
        height = 0.2  # 子图高度
        ax = plt.axes([left, bottom, width, height])
        ax.plot(single_sided_freqs, single_sided_fft_db, label='Frequency Domain (dB)')
        ax.set_xlabel('Frequency (Hz)')
        ax.set_ylabel('Magnitude (dB)')
        ax.set_title('Three - phase Motor Line Current Single - phase Waveform with Noise (Frequency Domain in dB)')
        ax.legend()
        ax.set_xlim(0, 2 * self.base_freq)
        ax.grid(True)

        plt.grid(True)

        plt.tight_layout()
        plt.show()


eca = EngineCurrentAnalyse1()
t,sig,fs = eca.GenCurrentData()
sig = eca.breakOneBar(t, sig)
eca.draw_signal(sig, t, fs)