import numpy as np
import math
import matplotlib.pyplot as plt
from matplotlib.widgets import RectangleSelector, SpanSelector



class SineGenerator:
    def __init__(self):
        self.curr_freq = 0.0
        self.coeff_A = 0.0
        self.coeff_B = 0.0
        self.state = [0.0, 0.0]
        self.PI = math.pi

    def set_freq(self, freq, sample_rate):
        """
        设置频率参数
        """
        if self.curr_freq != freq:
            # 二阶IIR振荡器
            number_in_period = sample_rate / freq
            self.coeff_A = math.sin(2 * self.PI / number_in_period * 2) / math.sin(2 * self.PI / number_in_period)
            self.coeff_B = -1.0
            self.state[0] = 0.0
            self.state[1] = math.sin(2 * self.PI / number_in_period)
            
            self.curr_freq = freq

    def generate_sine(self, number_samples):
        """
        生成正弦波样本
        返回numpy数组形式的音频缓冲区
        """
        audio_buffer = np.zeros(number_samples, dtype=np.float32)
        
        for i in range(number_samples):
            # IIR滤波器计算
            y = self.coeff_A * self.state[1] + self.coeff_B * self.state[0]
            
            # 更新状态
            self.state[0] = self.state[1]
            self.state[1] = y
            
            # 限制输出范围
            if y >= 1.0:
                y = 0.999999
            elif y < -1.0:
                y = -1.0
                
            audio_buffer[i] = y
        
        return audio_buffer

# 使用soundfile库的方法 (最现代的方法)
def save_float_to_wav_soundfile(data, filename, sample_rate=44100):
    """
    使用soundfile库将浮点数数组保存为WAV文件
    注意：需要先安装 soundfile: pip install soundfile
    """
    import soundfile as sf
    
    # soundfile可以直接处理浮点数，无需转换
    sf.write(filename, data, sample_rate)
# 使用示例
def create_sine_wave(freq, sample_rate, duration):
    """
    创建指定频率、采样率和持续时间的正弦波
    """
    generator = SineGenerator()
    generator.set_freq(freq, sample_rate)
    
    number_samples = int(sample_rate * duration)
    sine_wave = generator.generate_sine(number_samples)
    
    return sine_wave

# 示例：生成440Hz的音调，持续1秒，采样率44100Hz
sine_wave = create_sine_wave(440, 44100, 1.0)
x = np.arange(len(sine_wave))
fig, ax = plt.subplots()
line, = ax.plot(x, sine_wave, '-')
# SpanSelector 示例
span = SpanSelector(ax, onselect=lambda xmin, xmax: line.set_xdata(x[(x >= xmin) & (x <= xmax)]), 
                   direction='horizontal', useblit=True)

# RectangleSelector 示例
rect_selector = RectangleSelector(ax, onselect=lambda eclick, erelease: line.set_xdata(x[(x >= eclick[0]) & (x <= erelease[0])]), 
                                   useblit=True, button=[1, 3],  # 使用鼠标中键或右键拖动
                                  minspanx=5, minspany=5, spancoords='pixels')
plt.show()
save_float_to_wav_soundfile(sine_wave, 'sine_wave.wav')