from scipy.signal import chirp, spectrogram
import numpy as np
import matplotlib.pyplot as plt

def amp_freq_mod():
    # amplitude modulation
    t = np.linspace(0, 20 * np.pi, 1000, endpoint=False)
    t1 = t[0:200]
    t2 = t[200:600]
    t3 = t[600:800]
    t4 = t[800:]

    plt.figure(1)
    signal_unmod = np.sin(t)
    plt.plot(t, signal_unmod)
    plt.title("Unmodulated signal")
    plt.xlabel('time')
    plt.ylabel('Amplitude')
    plt.tight_layout()
    plt.savefig('unmodulated.png')
    plt.show()

    plt.figure(2)
    signal_ampmod = np.hstack((np.sin(t1), 0.2 * np.sin(t2), np.sin(t3), 0.2 * np.sin(t4)))
    plt.plot(t, signal_ampmod)
    plt.title("Amplitude modulation")
    plt.xlabel('time')
    plt.ylabel('Amplitude')
    plt.tight_layout()
    plt.savefig('amplitude_modulation.png')
    plt.show()

    # frequency modulation
    plt.figure(3)
    signal_ampmod = np.hstack((np.sin(t1), np.sin(0.5*t2), np.sin(t3), np.sin(0.5*t4)))
    plt.plot(t, signal_ampmod)
    plt.title("Frequency modulation")
    plt.xlabel('time')
    plt.ylabel('Amplitude')
    plt.tight_layout()
    plt.savefig('frequency_modulation.png')
    plt.show()

def chirps():
    # chirp example
    fs = 8000
    T = 10
    t = np.linspace(0, T, T*fs, endpoint=False)

    upchirp = chirp(t, f0=0, f1=5, t1=10, method='linear')
    plt.figure(1)
    plt.plot(t, upchirp)
    plt.title("Chirp")
    plt.xlabel('time')
    plt.ylabel('Amplitude')
    plt.tight_layout()
    plt.savefig('chirp.png')
    plt.show()


    #######################################
    # LoRa chirps
    fs = 8000
    T = 10
    t = np.linspace(0, T, T*fs, endpoint=False)

    upchirp = chirp(t, f0=250, f1=1750, t1=10, method='linear')

    ff, tt, Sxx = spectrogram(upchirp, fs=fs, noverlap=256, nperseg=512,
                                      nfft=2048)

    # Chirp explanation
    plt.figure(2)
    plt.yticks([250, 1000, 1750], ['f_start', 'f_center', 'f_end'] )
    plt.xticks([0,10], [0, 't_symb'])
    plt.pcolormesh(tt, ff[:513], Sxx[:513])
    plt.title("Chirp spectrogram")
    plt.xlabel('time')
    plt.ylabel('frequency')
    plt.tight_layout()
    plt.savefig('chirp_spectrogram.png')
    plt.show()

    # LoRa upchirp
    plt.figure(3)
    plt.yticks([250, 1000, 1750], ['868.1 MHz - 62.5 kHz', '868.1 MHz', '868.1 MHz + 62.5 kHz'] )
    plt.xticks([0,10], [0, 't_symb'])
    plt.pcolormesh(tt, ff[:513], Sxx[:513])
    plt.title("LoRa Up-Chirp")
    plt.xlabel('time')
    plt.ylabel('frequency')
    plt.tight_layout()
    plt.savefig('lora_upchirp.png')
    plt.show()


    # downchirp
    downchirp = chirp(t, f0=1750, f1=250, t1=10, method='linear')

    ff, tt, Sxx = spectrogram(downchirp, fs=fs, noverlap=256, nperseg=512,
                                      nfft=2048)


    plt.figure(4)
    plt.yticks([250, 1000, 1750], ['868.1 MHz - 62.5 kHz', '868.1 MHz', '868.1 MHz + 62.5 kHz'] )
    plt.xticks([0,10], [0, 't_symb'])
    plt.pcolormesh(tt, ff[:513], Sxx[:513])
    plt.title("LoRa Down-Chirp")
    plt.xlabel('time')
    plt.ylabel('frequency')
    plt.tight_layout()
    plt.savefig('lora_downchirp.png')
    plt.show()

def symbol():
    # symbol example
    fs = 8000
    T = 10
    t = np.linspace(0, T, T*fs, endpoint=False)

    fracs = [0.7, 0.7, 0.99, 0.9, 0.8, 0.5]
    for i in range(0,6):
        frac = fracs[i]
        t_jump = T * (1 - frac)

        # 250, 500, 750, 1000, 1250, 1500, 1750
        t1 = (frac)*T

        symbol_p1 = chirp(t[0:int(frac*len(t))], f0=250, f1=1750, t1=T, method='linear')
        symbol_p2 = chirp(t[int(frac*len(t)):], f0=250, f1=1750, t1=T, method='linear')
        symbol = np.hstack((symbol_p2, symbol_p1))
        ff, tt, Sxx = spectrogram(symbol, fs=fs, noverlap=256, nperseg=512,
                                          nfft=2048)

        #chirp = np.hstack((upchirp, upchirp))
        #ff, tt, Sxx = spectrogram(chirp, fs=fs, noverlap=256, nperseg=512,
        #                                  nfft=2048)

        plt.figure(5)
        plt.yticks([250, 1000, 1750], ['868.1 MHz - 62.5 kHz', '868.1 MHz', '868.1 MHz + 62.5 kHz'] )
        plt.pcolormesh(tt, ff[:513], Sxx[:513])
        if i > 0:
            plt.xticks([0, t_jump], [0, 't_jump'])
            plt.plot([t_jump, t_jump], [0, 2000.0], 'r')
        else:
            plt.xticks([0, t_jump], [0, 't_jump'])
            plt.gca().tick_params(axis='x', colors='white')
        plt.title("LoRa Symbol")
        plt.xlabel('time')
        plt.ylabel('frequency')
        plt.tight_layout()
        plt.savefig('lora_symbols_{}.png'.format(i))
        plt.show()


def spreading_factor():
    fs = 8000
    T = 400
    t = np.linspace(0, T, T * fs, endpoint=False)

    chirps = np.array([0])
    for i in range(3,-1,-1):
        t_final = len(t)/2**i
        chirp1 = chirp(t[0:t_final], f0=250, f1=1750, t1=T/2**i, method='linear')

        chirps = np.hstack((chirps, chirp1))

    ff, tt, Sxx = spectrogram(chirps, fs=fs, noverlap=256, nperseg=512,
                              nfft=2048)

    plt.figure(5)
    plt.yticks([250, 1000, 1750], ['868.1 MHz - 62.5 kHz', '868.1 MHz', '868.1 MHz + 62.5 kHz'])
    plt.pcolormesh(tt, ff[:513], Sxx[:513])

    plt.title("LoRa Spreading Factors")
    plt.xlabel('time')
    plt.ylabel('frequency')
    plt.xticks([], [])
    plt.gca().set_aspect(0.28)

    for i in range(0,4):
        plt.text(tt[int(2.0**i/2.0**3 * len(tt)-1)], 1800.0, 'SF = {}'.format(i + 7), color='white', horizontalalignment='right')

    plt.tight_layout()
    plt.savefig('lora_spreading_factors.png', dpi=400.0)
    plt.show()
    pass


#symbol()
spreading_factor()

plt.show()