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


def wgn(x, snr=10):
    """ 给数据 x 添加指定信噪比为 SNR dB 的高斯噪声 """
    np.random.seed(100)
    snr = 10 ** (snr / 10.0)
    xpower = np.sum(x ** 2) / len(x)
    npower = xpower / snr
    return np.random.randn(len(x)) * np.sqrt(npower)

def gener_data():
    t = np.linspace(0, 2, num=2000)
    x1 = 500 * np.sin(2 * np.pi * t)
    x2 = 1000 * np.sin(2 * np.pi * t)
    x3 = 2000 * np.sin(2 * np.pi * t)
    x4 = 5000 * np.sin(2 * np.pi * t)

    noise_1 = wgn(x1, snr=20)
    noise_2 = wgn(x2, snr=15)
    noise_3 = wgn(x3, snr=10)
    noise_4 = wgn(x4, snr=5)

    x1 += noise_1
    x2 += noise_2
    x3 += noise_3
    x4 += noise_4

    x1 = x1 + np.array([0] * 1000 + [750] * 1000)
    x2 = x2 + np.array([0] * 1000 + [1500] * 1000)
    x3 = x3 + np.array([0] * 1000 + [3000] * 1000)
    x4 = x4 + np.array([0] * 1000 + [7500] * 1000)
    print(x1.shape)

    return x1, x2, x3, x4

def _plot(x1, x2, x3, x4):
    plt.figure(figsize=(16, 8))
    plt.plot(x1)
    plt.plot(x2)
    plt.plot(x3)
    plt.plot(x4)
    plt.tight_layout()
    plt.show()

def write_data(x1, x2, x3, x4):
    df = pd.DataFrame({
   "x1":x1,
                       "x2":x2,
                       "x3":x3,
                       "x4":x4})
    df.to_csv("./simu_data.csv", index=False, sep=",")

if __name__ == "__main__":
    x1, x2, x3, x4 = gener_data()
    _plot(x1, x2, x3, x4)
    write_data(x1, x2, x3, x4)