# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2023/1/1
"""
绘制 窗函数和对应的频率响应
"""
import numpy as np
from numpy.fft import rfft
import matplotlib.pyplot as plt
window_len = 60
 
 
# frequency response
def frequency_response(window, window_len=window_len, NFFT=2048):
    A = rfft(window, NFFT) / (window_len / 2)  # (513,)
    mag = np.abs(A)
    freq = np.linspace(0, 0.5, len(A))
    # 忽略警告
    with np.errstate(divide='ignore', invalid='ignore'):
        response = 20 * np.log10(mag)
    response = np.clip(response, -150, 150)
    return freq, response
def Rectangle_windows(win_length):
    # 矩形窗
    return np.ones((win_length))
def Voibis_windows(win_length):
    """ Voibis_windows窗函数，RNNoise使用的是它，它满足Princen-Bradley准则。
    :param x:
    :param win_length: 窗长
    :return:
    """
    x = np.arange(0, win_length)
    return np.sin((np.pi / 2) * np.sin((np.pi * x) / win_length) ** 2)
def sqrt_hanning_windows(win_length, mode="periodic"):
    # symmetric: 对称窗，主要用于滤波器的设计
    # periodic: 周期窗，常用于频谱分析
    if mode == "symmetric":
        haning_window = np.hanning(win_length)
        sqrt_haning_window = np.sqrt(haning_window)
    elif mode == "periodic":
        haning_window = np.hanning(win_length+1)
        sqrt_haning_window = np.sqrt(haning_window)
        sqrt_haning_window = sqrt_haning_window[0:-1].astype('float32')
    return sqrt_haning_window
Rectangle_windows = Rectangle_windows(window_len)
hanning_window = np.hanning(M=window_len)
print(np.argmax(hanning_window))
sqrt_hanning_windows = sqrt_hanning_windows(window_len)
hamming_window = np.hamming(M=window_len)
Voibis_windows = Voibis_windows(window_len)
blackman_window = np.blackman(M=window_len)
bartlett_window = np.bartlett(M=window_len)
kaiser_window = np.kaiser(M=window_len, beta=14)
plt.figure()
plt.plot(Rectangle_windows, label="Rectangle")
plt.plot(hanning_window, label="hanning")
plt.plot(sqrt_hanning_windows, label="sqrt_hanning")
plt.plot(hamming_window, label="hamming")
plt.plot(Voibis_windows, label="Voibis")
plt.plot(blackman_window, label="blackman")
plt.plot(bartlett_window, label="bartlett")
plt.plot(kaiser_window, label="kaiser")
plt.legend()
plt.tight_layout()
plt.show()
freq, Rectangle_FreqResp = frequency_response(Rectangle_windows, window_len)
freq, hanning_FreqResp = frequency_response(hanning_window, window_len)
freq, sqrt_hanning_FreqResp = frequency_response(sqrt_hanning_windows, window_len)
freq, hamming_FreqResp = frequency_response(hamming_window, window_len)
freq, Voibis_FreqResp = frequency_response(Voibis_windows, window_len)
freq, blackman_FreqResp = frequency_response(blackman_window, window_len)
freq, bartlett_FreqResp = frequency_response(bartlett_window, window_len)
freq, kaiser_FreqRespw = frequency_response(kaiser_window, window_len)
plt.figure()
plt.title("Frequency response")
plt.plot(freq, Rectangle_FreqResp, label="Rectangle")
plt.plot(freq, hanning_FreqResp, label="hanning")
plt.plot(freq, sqrt_hanning_FreqResp, label="sqrt_hanning")
plt.plot(freq, hamming_FreqResp, label="hamming")
plt.plot(freq, Voibis_FreqResp, label="Voibis")
plt.plot(freq, blackman_FreqResp, label="blackman")
plt.plot(freq, bartlett_FreqResp, label="bartlett")
plt.plot(freq, kaiser_FreqRespw, label="kaiser")
plt.ylabel("Magnitude [dB]")
plt.xlabel("Normalized frequency [cycles per sample]")
plt.legend()
plt.tight_layout()
plt.show()

# -*- coding:utf-8 -*-  
# Author:凌逆战 | Never# Date: 2024/3/8  
"""  
对比不同库中hann窗函数的实现  
如果对称(sym=True)的话，有两个最大值，如果不对称(sym=False)的话，有一个最大值  
  
- numpy的hanning函数是对称的  - scipy有hanning函数有sym参数设置，默认是对称的  
- torch的hanning函数有periodic参数设置，默认是非对称的  
"""  
import numpy as np  
  
import torch  
import scipy.signal as signal  
  
window_len = 512  
  
  
def hann_sym(window_len):  
    """对称hann窗"""  
    win = np.zeros(window_len)  
    for i in range(window_len):  
        win[i] = 0.5 - 0.5 * np.cos(2 * np.pi * i / (window_len - 1))  
    return win  
  
  
def hann_asym(window_len):  
    """非对称hann"""  
    p_win = np.zeros(window_len)  
    for i in range(window_len):  
        p_win[i] = np.sin(np.pi * i / window_len)  
        p_win[i] = p_win[i] * p_win[i]  
    return p_win  
  
  
def my_hann_aysm(win_len):  
    haning_window = np.hanning(win_len + 1)  # 对称的hann窗  
    out = haning_window[0:-1].astype('float32')  # 舍弃最后一个元素  
    return out  
  
  
scipy_sym = signal.windows.hann(window_len, sym=True)  # 对称的hann窗  
scipy_Asym = signal.windows.hann(window_len, sym=False)  # 非对称的hann窗  
hann_sym_c = hann_sym(window_len)  
hann_asym_c = hann_asym(window_len)  
my_hann= my_hann_aysm(window_len)  
  
print(np.allclose(scipy_sym, hann_sym_c))  # True  
print(np.allclose(scipy_Asym, hann_asym_c))  # True  
print(np.allclose(my_hann, hann_asym_c))  # True  
  
numpy_window = np.hanning(window_len)  # 说明numpy的hanning函数是对称的  
print(np.allclose(numpy_window, scipy_sym))  # True  
  
torch_window = torch.hann_window(window_len)  # 非对称  
torch_window_periodic = torch.hann_window(window_len, periodic=False)  # 非周期=对称  
# print(torch.argmax(window_torch))  
  
# 判断两个窗函数是否相等  
print(np.allclose(scipy_Asym, torch_window.numpy(), rtol=1e-3))  # True  
print(np.allclose(scipy_sym, torch_window_periodic.numpy(), rtol=1e-3))  # True

def Orka_forward_window(N1=64, N2=448, hop_size=64, NFFT=512):  
    analysisWindow = np.zeros(NFFT)  
    for n in range(NFFT):  
        if n < N1:  
            analysisWindow[n] = np.sin(n * np.pi / (2 * N1)) ** 2  
        elif N1 <= n <= N2:  
            analysisWindow[n] = 1  
        elif N2 < n <= N2 + hop_size:  
            analysisWindow[n] = np.sin(np.pi * (N2 + hop_size - n) / (2 * hop_size))  
  
    return analysisWindow  
  
  
def Orka_backward_window(N1=64, N2=448, hop_size=64, NFFT=512):  
    synthesisWindow = np.zeros(NFFT)  
    for n in range(NFFT):  
        if n < N2 - hop_size:  
            synthesisWindow[n] = 0  
        elif N2 - hop_size <= n <= N2:  
            synthesisWindow[n] = np.cos(np.pi * (n - N2) / (2 * hop_size)) ** 2  
        elif N2 < n <= N2 + hop_size:  
            synthesisWindow[n] = np.sin(np.pi * (N2 + hop_size - n) / (2 * hop_size))  
    return synthesisWindow

def TukeyAW(n, N, alpha):  
    # assert n >= 0  
    if n < alpha * N:  
        return 0.5 * (1 - np.cos(np.pi * n / (alpha * N)))  
    elif n <= N - alpha * N:  
        return 1  
    elif n <= N:  
        return 0.5 * (1 - np.cos(np.pi * (N - n) / (alpha * N)))  
  
  
def getTukeyAnalysisWindow(filter_length, alpha):  
    analysisWindow = np.zeros(filter_length)  
    for i in range(filter_length):  
        analysisWindow[i] = TukeyAW(i, filter_length, alpha)  
    return analysisWindow  
  
  
def getTukeySynthesisWindow(N, A, B, alpha):  
    synthesisWindow = np.zeros(A)  
    for i in range(A):  
        x = N - A + i  
        numerator = TukeyAW(x, N, alpha)  
        denonminator = 0  
        for k in range(int(A / B)):  
            y = N - A + i % B + k * B  
            denonminator += TukeyAW(y, N, alpha) ** 2  
        synthesisWindow[i] = numerator / denonminator  
  
    synthesisWindow = np.pad(synthesisWindow, (N - A, 0), 'constant', constant_values=0)  
    return synthesisWindow  

def getAsqrtAnalysisWindow(N, M, d):  
    # filter_length, hop_length, d  
    risingSqrtHann = np.sqrt(np.hanning(2 * (N - M - d) + 1)[:(N - M - d)])  
    fallingSqrtHann = np.sqrt(np.hanning(2 * M + 1)[:2 * M])  # 下降  
  
    window = np.zeros(N)  
    window[:d] = 0  
    window[d:N - M] = risingSqrtHann[:N - M - d]  
    window[N - M:] = fallingSqrtHann[-M:]  
  
    return window  
  
  
def getAsqrtSynthesisWindow(N, M, d):  
    risingSqrtHannAnalysis = np.sqrt(np.hanning(2 * (N - M - d) + 1)[:(N - M - d)])  
    fallingSqrtHann = np.sqrt(np.hanning(2 * M + 1)[:2 * M])  
    risingNoramlizedHann = np.hanning(2 * M + 1)[:M] / risingSqrtHannAnalysis[N - 2 * M - d:N - M - d]  
  
    window = np.zeros(N)  
    window[:-2 * M] = 0  
    window[-2 * M:-M] = risingNoramlizedHann  
    window[-M:] = fallingSqrtHann[-M:]  
  
    return window