import numpy as np
import scipy.io.wavfile as wf
import scipy.signal as sig
import sounddevice as sd

# 生成sin,cos信号
def sin_signal():
    t = np.arange(0, 1.01, 1/100)
    sin=np.sin(2*np.pi*1*t)
    cos=np.cos(2*np.pi*1*t)
    return sin,cos

# iq信号合成
def iq_modulation(FS=100,Fi=10,Fq=10,T=2):
    t = np.arange(0, T, 1/FS)
    i=np.cos(2*np.pi*Fi*t)
    q=np.sin(2*np.pi*Fq*t)
    iq=i+q
    return i,q,iq

# am调制
def am_modulation(FS=100,Fm=1,Fc=10,T=2,m=0.7):
    # AM=A*(1+Ma*Wm*t)*cos(Wc*t)
    t = np.arange(0, T, 1/FS)
    modulator = np.cos(2*np.pi*Fm*t)
    carrier = np.cos(2*np.pi*Fc*t)
    product = (1+m*modulator)*carrier
    return modulator,carrier,product

# fm调制
def fm_modulation(FS=100,Fm=1,Fc=10,T=2,m=5):
    # FM=A*cos(Wc*t+Ma*sin(Wm*t))
    t = np.arange(0, T, 1/FS)
    modulator=np.cos(2*np.pi*Fm*t)
    carrier=np.cos(2*np.pi*Fc*t)
    product = np.cos(2*np.pi*Fc*t+m*np.sin(2*np.pi*Fm*t))
    return modulator,carrier,product

# iq文件读取
def read_iq_file(filename):
    fs,rf = wf.read(filename)
    print("sampling rate = {} Hz, length = {} samples, channels = {}, dtype = {}".format(fs, *rf.shape, rf.dtype))
    return fs,(rf[:, 0] + 1j * rf[:, 1])

# 频率偏移
def frequency_shift(data,fs,offset):
    return data * np.exp(1.0j*2.0*np.pi* offset/fs*np.arange(len(data)))

# 重采样
def resample_poly(data,up,down):
    return sig.resample_poly(data, up,down)     

# 移除直流-平均法
def remove_dc(signal):
    return signal-np.mean(signal)

# 移除直流-过滤器
def remove_dc_filter(signal, fs, cutoff_hz):
    b, a = sig.butter(1, cutoff_hz / (fs/2), 'highpass')
    return sig.lfilter(b, a, signal)

# am解调
def am_demodulation(data):
    # sqrt(I*I+Q*Q)=A0+m(n)
    # return np.sqrt(data.real**2+data.imag**2)
    return np.abs(data)

# fm解调
def fm_demodulation(data):
    # φ(n)-φ(n-1)=m(n)
    # f(n)=I(n-1)*Q-I*Q(n-1)
    # return np.angle(data[1:] * np.conj(data[:-1]))  
    return np.diff(np.unwrap(np.angle(data)))

# 去加重过滤器
def lfilter_data(data,bw):
    d = bw * 75e-6 
    x = np.exp(-1/d)
    b = [1-x]
    a = [1,-x]  
    ldata = sig.lfilter(b,a,data) 
    return ldata

# 音频适应缩放
def scale_to_volume(data,scale=32768):
    return data*scale / np.max(np.abs(data)) 

# 音频文件写入
def write_to_file(data,filename):
    wf.write(filename,44100,data.astype("int16"))

# 音频播放
def play_audio(data,fs):
    sd.play(data, fs)
    sd.wait()

# 带通采样计算:fs=4*f0/(2*n+1),fs>=2*bw
def computeFs(f0,bw):
    n=1
    while(True):
        fs=4*f0/(2*n+1)
        if(fs<=2*bw):
            return
        if(fs%1==0):
            print(fs)
        n+=1

__all__=[
    'sin_signal',
    'iq_modulation',
    'read_iq_file',
    'frequency_shift',
    'resample_poly',
    'am_modulation',
    'am_demodulation',
    'fm_modulation',
    'fm_demodulation',
    'lfilter_data',
    'scale_to_volume',
    'write_to_file',
    'play_audio',
    'computeFs',
]