import numpy as np
import scipy.signal as sig
import matplotlib.pyplot as plt
from sdr import *

def test_iq_modulation():
    i,q,iq=iq_modulation(FS=100,Fi=2,Fq=2,T=2)
    plt.plot(i,'r',label='i')
    plt.plot(q,'g',label='q')
    plt.plot(iq,'b',label='i+q')
    plt.legend()
    plt.show()

def test_am_modulation():
    modulator,carrier,product=am_modulation(FS=100,Fm=1,Fc=10,T=2,m=0.7)
    plt.title('AM Modulation')
    plt.subplot(3,1,1)
    plt.title('modulator')
    plt.plot(modulator,'r')
    plt.subplot(3,1,2)
    plt.title('carrier')
    plt.plot(carrier,'g')
    plt.subplot(3,1,3)
    plt.title('product')
    plt.plot(product,'b')
    demod=am_demodulation(product)
    plt.subplots_adjust(hspace=1)
    plt.show()

def test_fm_modulation():
    modulator,carrier,product=fm_modulation(FS=100,Fm=1,Fc=10,T=2,m=5)
    plt.title('FM Modulation')
    plt.subplot(3,1,1)
    plt.title('modulator')
    plt.plot(modulator,'r')
    plt.subplot(3,1,2)
    plt.title('carrier')
    plt.plot(carrier,'g')
    plt.subplot(3,1,3)
    plt.title('product')
    plt.plot(product,'b')
    plt.subplots_adjust(hspace=1)
    plt.show()

def test_am_demodulation():
    fs,iq=read_iq_file('./iq-files/SDRSharp_20230730_073925Z_0Hz_IQ.wav')
    shift_data=frequency_shift(iq,fs,549*1000)#550，-900
    bw_data=resample_poly(shift_data,20,3000)
    angle_data=am_demodulation(bw_data)
    ldata=lfilter_data(angle_data,200e3)
    audio_data=resample_poly(ldata,441,200)
    volume_data=scale_to_volume(audio_data,18000)
    write_to_file(volume_data,'./wav-files/am.wav')

    # plt.psd(iq,Fs=fs)
    # plt.subplots_adjust(hspace=1)
    # plt.show()

def test_fm_demodulation():
    fs,iq=read_iq_file('./iq-files/SDRSharp_20230729_104201Z_98760kHz_IQ.wav')
    shift_data=frequency_shift(iq,fs,540*1000)
    bw_data=resample_poly(shift_data,200,2000)
    angle_data=fm_demodulation(bw_data)
    # plt.plot(angle_data[0:1000],'r',label='angle_data')
    ldata=lfilter_data(angle_data,200e3)
    # plt.plot(ldata[0:1000],'g',label='ldata')
    audio_data=resample_poly(ldata,441,2000)
    volume_data=scale_to_volume(audio_data)
    write_to_file(volume_data,'./wav-files/fm.wav')

    # plt.psd(shift_data,Fs=fs)
    # plt.subplots_adjust(hspace=1)
    # plt.legend()
    # plt.show()

def test_frequency_shift():
    fs=100
    _,_,product=am_modulation(FS=fs,Fm=1,Fc=10,T=2,m=0.7)

    plt.subplot(3,1,1)
    plt.title('product')
    plt.plot(product,'r')

    plt.subplot(3,1,2)
    plt.title('product')
    plt.psd(product,Fs=fs)

    shift_data=frequency_shift(product,fs,-10)
    plt.subplot(3,1,3)
    plt.title('shift-data')
    plt.psd(shift_data,Fs=fs)

    plt.subplots_adjust(hspace=1)
    plt.show()

def test_sin_signal():
    sin,cos=sin_signal()
    plt.plot(sin,'r')
    plt.plot(1.1*sin,'g')
    plt.legend()
    plt.grid()
    plt.show()

def test_exp_cos():
    t = np.arange(0, 1,0.001)
    x=np.exp(2j*np.pi*1*t)
    x2=np.cos(2*np.pi*1*t)

    plt.subplot(2,1,1)
    plt.title('np.exp(2j*np.pi*1*t)')
    plt.grid()
    plt.plot(x,'r')
    plt.subplot(2,1,2)
    plt.title('np.cos(2*np.pi*1*t)')
    plt.grid()
    plt.plot(x2,'g')
    plt.subplots_adjust(hspace=1)
    plt.show()
    return

def test_resample_poly():
    t = np.arange(0, 1,0.01)
    x=np.exp(2j*np.pi*1*t)
    x2=sig.resample_poly(x, 1,5)
        
    plt.subplot(2,1,1)
    plt.title('x')
    plt.grid()
    plt.plot(x,'o-r')
    plt.subplot(2,1,2)
    plt.title('resample_poly(x,1,10)')
    plt.grid()
    plt.plot(x2,'o-g')
    plt.subplots_adjust(hspace=1)
    plt.show()

def test_computeFs():
    computeFs(450e3,10e3)

test_computeFs()