# -*- coding:utf-8 -*-
# Create by Zhou
# 2021/11/29/17:02
# May Saint Diana bless your coding!

import numpy as np
import matplotlib.pyplot as plt
from scipy import signal


def insertToHead(x: np.array, element):
    """
    将x末尾删去，并把element加入到头部
    :param x:
    :param element:
    :return: 返回更新后的x
    """
    n = len(x)
    x = np.delete(x, n - 1)  # 删除最后一个元素
    x = np.insert(x, 0, element)  # 新元素加入到头部
    return x


def insertToTail(x: np.array, element):
    """
    将x的头部删去，并将element加到末尾
    :param x: 待更新的array
    :param element: 要加入的元素
    :return: 返回更新后的x
    """
    x = np.append(x, element)  # 插入到末尾
    x = np.delete(x, 0)  # 删除头部元素
    return x


def FIR_Filter(inputSignal: np.array, model: np.array) -> np.array:
    """
    FIR线性滤波器,得到经过输入信号经过滤波器滤波的结果
    y(n) = b(1)*x(n) + b(2)*x(n-1) + ... + b(nb+1)*x(n-nb)
    :param inputSignal: 输入时序信号
    :param model: 输入FIR模型
    :return: 返回滤波结果 长度和输入相等
    """
    m = len(inputSignal)
    n = len(model)
    assert m >= n
    output = np.zeros(m)
    window = np.zeros(n)  # 计算窗  和模型系数进行点乘
    for i in range(m):
        window = insertToHead(window, inputSignal[i])
        output[i] = np.dot(window, model)  # 点乘 对应位相乘后求和
    return output


def FIR_Identify(x, d, length, u):
    """
    利用FIR滤波器，得到x和d之间的线性关系
    :param x: 输入
    :param d: 输出
    :param length: FIR滤波器长度
    :param u: 迭代过程的步长因子
    :return: 返回最终的FIR滤波器系数及整个过程的拟合误差
    """
    n = len(x)
    w = np.zeros(length)  # FIR滤波器系数
    cx = np.zeros(length)  # 输入信号及其延迟
    error = np.zeros(n)
    for i in range(n):
        # 更新输入
        cx = insertToHead(cx, x[i])
        # 计算误差
        error[i] = d[i] - np.dot(cx, w)
        # 更新系数
        w = w + u * error[i] * cx
    return w, error


def FxLMS(x, Yp, N, secModel, u):
    """
    FxLMS算法
    :param x: 输入参考信号
    :param Yp: 参考信号经过初级路径得到的初级噪声，也就是我们要消除的
    :param N:  FIR滤波器的长度，也就是系数W的长度
    :param secModel: 次级路径FIR模型
    :param u: 步长因子
    :return: 返回拟合误差
    """
    Cx = np.zeros(N)  # 控制器处得到的当前信号x(n)及其延迟信号
    W = np.zeros(N)  # 控制器（本质上是一个FIR滤波器）权重
    Yw = np.zeros(len(secModel))  # 次级路径处得到的控制器信号Yw(k)及其延迟
    T = len(x)  # % 工作时间，和X的长度一致
    error = np.zeros(T)  # % 误差传感器处的误差信号
    Fx = np.zeros(N)  # % 滤波x信号
    for i in range(T):
        #  计算输出
        Cx = insertToHead(Cx, x[i])  # 更新x(n)信号
        Cy = np.dot(Cx, W)  # 控制器根据当前信号及其延迟信号，计算输出
        Yw = insertToHead(Yw, Cy)  # 控制器输出的信号传播到次级路径，进入延迟队列
        #  计算误差
        error[i] = Yp[i] + np.dot(Yw, secModel)  # % 输出到次级扬声器的信号经过次级路径到达测量点，和初级信号叠加，构成误差信号
        #  更新系数
        Fx = insertToHead(Fx, np.dot(Cx[0:len(secModel)], secModel))  # 计算经过次级通道估计器过滤得到滤波x信号
        W = W - u * error[i] * Fx  # 更新FIR控制器滤波系数
    return error


def ANR(desire, error, u=0.999):
    """
    average noise reduction
    ANR(n) = 20log(Ae(n) / Ad(n))
    Ae(n) = uAe(n - 1) + (1 - u) | e(n) |
    Ad(n) = uAd(n - 1) + (1 - u) | d(n) |
    u is the forgetting factor set  to 0.999
    :param desire: 拟合目标
    :param error:  拟合误差
    :param u: 遗忘因子
    :return:  ANR 平均降噪量
    """
    T = len(desire)
    Ae = np.ones(T)
    Ad = np.ones(T)
    for t in range(1, T):
        Ae[t] = max(u * Ae[t - 1] + (1 - u) * np.abs(error[t]), 0.001)
        Ad[t] = max(u * Ad[t - 1] + (1 - u) * np.abs(desire[t]), 0.001)
    y = 20 * np.log10(Ae / Ad)
    return y


def MSE(before, after):
    """
    均方误差  用来评价降噪量 (dB)
    :param after:
    :param before:
    :return:
    """
    mse = np.abs((before - after) / before)
    mse = 20 * np.log10(mse)  # 幅值20log 能量10log
    return mse


def sin_wave(A, f, fs, phi, t):
    """
    :params A:    振幅
    :params f:    信号频率
    :params fs:   采样频率
    :params phi:  相位
    :params t:    时间长度
    """
    # 若时间序列长度为 t=1s,
    # 采样频率 fs=1000 Hz, 则采样时间间隔 Ts=1/fs=0.001s
    # 对于时间序列采样点个数为 n=t/Ts=1/0.001=1000, 即有1000个点,每个点间隔为 Ts
    Ts = 1 / fs
    n = t / Ts
    n = np.arange(n)
    y = A * np.sin(2 * np.pi * f * n * Ts + phi * (np.pi / 180))
    return y


def awgn(x, snr):
    """
    给x添加一定分贝的高斯白噪声
    :param x:
    :param snr: 信噪比  snr越大 噪声越小
    :return:
    """
    Ps = np.sum(abs(x) ** 2) / len(x)  # 原始信号功率
    Pn = Ps / (10 ** (snr / 10))  # 噪声功率
    noise = np.random.randn(len(x)) * np.sqrt(Pn)  # 噪声
    signal_add_noise = x + noise[:, np.newaxis]
    return signal_add_noise


def NPP_1(x):
    """
    nonlinear primary path model 非线性初级路径模型1
    p(n) = 0.4x(n-3) - 0.3x(n-4) + 0.2x(n-5) + 0.2x(n-5)^2
    :param x:
    :return:
    """
    Px = np.zeros(6)
    Yp = np.zeros_like(x)
    T = len(x)
    for i in range(T):
        Px = insertToHead(Px, x[i])
        Yp[i] = 0.4 * Px[3] - 0.3 * Px[4] + 0.2 * Px[5] + 0.2 * Px[5] ** 2
    return Yp


def NPP_2(x):
    """
    nonlinear primary path model 非线性初级路径模型2
    % FsLMS论文用
    d(n) = u(n-2) + 0.08u(n-2)^2 - 0.04u(n-1)^3
    u(n) = x(n) * f(n)
    f(n) 对应 [0 0 0 1 -0.3 0.2]
    :param x:
    :return:
    """
    f = np.array([0, 0, 0, 1, - 0.3, 0.2])
    Fx = np.zeros(6)
    T = len(x)
    Yp = np.zeros(T)
    U = np.zeros(3)
    for i in range(T):
        Fx = insertToHead(Fx, x[i])
        U = insertToHead(U, np.dot(Fx, f))
        Yp[i] = U[2] + 0.08 * U[2] ** 2 - 0.04 * U[1] ** 3
    return Yp


def NPP_3(x):
    """
    nonlinear primary path model 非线性初级路径模型3
    GFLANN论文  case1 初级路径
    d(n) = x(n)+0.8x(n-1) +0.3x(n-2) + 0.4x(n-3) - 0.8x(n)x(n-1) +0.9x(n)x(n-2) +0.7x(n)x(n-3)
    :param x:
    :return:
    """
    T = len(x)
    Px = np.zeros(4)
    Yp = np.zeros(T)
    for i in range(T):
        Px = insertToHead(Px, x[i])
        Yp[i] = Px[0] + 0.8 * Px[1] + 0.3 * Px[2] + 0.4 * Px[3] - 0.8 * Px[0] * Px[1] + \
                0.9 * Px[0] * Px[2] + 0.7 * Px[0] * Px[3]
    return Yp


def NPP_4(x):
    """
    nonlinear primary path model 非线性初级路径模型4
    GFLANN论文  case2 初级路径
    d(n) = x(n)+0.8x(n-1) +0.3x(n-2) + 0.4x(n-3) - 0.8x(n)*x(n-1) +0.9x(n)*x(n-2)
            +0.7x(n)*x(n-3) -3.9x(n-1)^2*x(n-2)-2.6x(n-1)^2*x(n-3)+2.1x(n-2)^2*x(n-3)
    太复杂了 改一下：
        d(n) = x(n) + 0.8x(n-1) - 0.4x(n-3) - 0.8x(n)*x(n-1) +0.9x(n)*x(n-2) + 0.6x(n-1)^2*x(n-3)
    :param x:
    :return:
    """
    T = len(x)
    Px = np.zeros(4)
    Yp = np.zeros(T)
    for i in range(T):
        Px = insertToHead(Px, x[i])
        Yp[i] = Px[0] + 0.8 * Px[1] - 0.4 * Px[3] - 0.8 * Px[0] * Px[1] + \
                0.9 * Px[0] * Px[2] + 0.6 * (Px[1] ** 2) * Px[3]
    return Yp


def NPP_5(x):
    """
    nonlinear primary path model 非线性初级路径模型5
    FLANN convex Volterra 论文用
    d(n) = x(n-5)+0.8x(n-6)+0.3x(n-7)+0.4x(n-8)+0.2x(n-5)x(n-6) -0.3x(n-5)x(n-7)+0.4x(n-5)x(n-8)
    :param x:
    :return:
    """
    T = len(x)
    Px = np.zeros(9)
    Yp = np.zeros(T)
    for i in range(T):
        Px = insertToHead(Px, x[i])
        Yp[i] = Px[5] + 0.8 * Px[6] + 0.3 * Px[7] + 0.4 * Px[8] + 0.2 * Px[5] * Px[6] \
                - 0.3 * Px[5] * Px[7] + 0.4 * Px[5] * Px[8]
    return Yp


def logisticChaotic(x0, r, length):
    noise = np.zeros(length)
    noise[0] = x0
    for i in range(1, length):
        noise[i] = r * noise[i - 1] * (1 - noise[i - 1])
    return noise


def plotSTFT(x0, fs, nPerSeg=256):
    f, t, Zxx = signal.stft(x0, fs, nperseg=nPerSeg)
    plt.figure()
    plt.pcolormesh(t, f, np.abs(Zxx),shading='auto')
    plt.title('STFT Magnitude')
    plt.ylabel('Frequency [Hz]')
    plt.xlabel('Time [sec]')
    plt.show()


if __name__ == '__main__':
    # x = sin_wave(A=5, f=50 + 450 * np.random.random(), fs=20000, phi=np.random.random(), t=1)
    fs = 10e3
    N = 1e5
    amp = 2 * np.sqrt(2)
    noise_power = 0.01 * fs / 2
    time = np.arange(N) / float(fs)
    mod = 500 * np.cos(2 * np.pi * 0.25 * time)
    carrier = amp * np.sin(2 * np.pi * 3e3 * time + mod)
    noise = np.random.normal(scale=np.sqrt(noise_power), size=time.shape)
    noise *= np.exp(-time / 5)
    on_carrier = carrier + noise
    plotSTFT(on_carrier, fs, 1000)