## Imports ##
from rtlsdr import RtlSdr
import numpy
import matplotlib.pyplot as plt
import time
import peakutils
import csv
import sys
import gpsd
#from gps3 import gps3
#import schedule



try:
    import RPi.GPIO as gpio
    
except:
    print("GPIO Pins Not Available on This Device\n")
 
 
## Function Definitions ##

# def gps_loc():
#             
#     # Attempt to connect to GPS Dongle
#     try:
#         gps_socket = gps3.GPSDSocket()
#         data_stream = gps3.DataStream()
#         
#         gps_socket.connect()
#         gps_socket.watch()
#         
#         # create counter to monitor number of attempts to find location
#         loc_count = 0
#         new_count = 0
#         
#         # wait for new data
#         for new_data in gps_socket:
#             if new_data:
#                 data_stream.unpack(new_data)
#                 
#                 # Find fix mode of device. 2 or 3 indicate lock. 1 is no lock. n/a is unknown     
#                 loc_mode = data_stream.TPV['mode']
#                 
#                 # if fix mode is known
#                 if loc_mode != 'n/a':
#                     #change mode from string to integer
#                     loc_mode = int(loc_mode)
#                     
#                     # If at least 2D fix (mode 2) is achieved, save longitude and latitude with dummy variable for altitude
#                     if loc_mode == 2:
#                         lon = data_stream.TPV['lon']
#                         lat = data_stream.TPV['lat']
#                         alt = 0
#                         break
#                      
#                     # If at 3D fix (mode 3) is achieved, save longitude, latitude and altitude
#                     elif loc_mode == 3:
#                         lon = data_stream.TPV['lon']
#                         lat = data_stream.TPV['lat']
#                         alt = data_stream.TPV['alt']
#                         break
#                     
#                     else:
#                         loc_count += 1
#                     
#                 # If device takes too long to find location, terminate.
#                 elif loc_count > 10:
#                     lon = 0
#                     lat = 0
#                     alt = 0
#                     break
#                 
#                 #advance attempts counter        
#                 else:
#                     loc_count += 1
#         
#             elif new_count > 1000000:
#                 lon = 0
#                 lat = 0
#                 alt = 0
#                 break
#             
#             else:
#                 new_count += 1
#         
#         #disconnect from GPS
#         gps_socket.close()   
#                 
#     # if unable to connect, save dummy variables, alert user and exit
#     except:
#         print('GPS Not Available on This Device\n')
#         lon = 0
#         lat = 0
#         alt = 0
#     
#     return lon, lat, alt

# # Function to read GPS location from a USB GPS dongle
def gps_loc():
       
    # Attempt to connect to GPS Dongle
    try:
        # Connect to the local gpsd
        #gpsd.connect()
   
        # Get gps position
        location = gpsd.get_current()
   
        # If at least 2D fix (mode 2 or 3) is achieved, save longitude and latitude
        if location.mode >= 2:
            lon = location.lon
            lat = location.lat
        else:
            lon = 0
            lat = 0
           
        # If at 3D fix (mode 3) is achieved, save altitude
        if location.mode >= 3:
            alt = location.alt
        else:
            alt = 0
       
    # if unable to connect, save dummy variables, alert user and exit
    except:
        print('GPS Not Available\n')
        lon = 0
        lat = 0
        alt = 0
       
    return lon, lat, alt

# Function to produce strings with time and date
def time_date():
    # string including date and time
    current_time = time.strftime("%d/%m/%y %H:%M:%S")
    # string with just date
    current_date = time.strftime("%d/%m/%y")
    # string with just time
    current_hour = time.strftime("%H:%M:%S")
    # string with date and time in a format suitable for use as a file name
    time_for_save = time.strftime("%m_%d_%y-%H_%M_%S")
    
    return current_date, current_hour, current_time, time_for_save

# Function to detect peaks in the data
def peak_det(spectrum, faxis, freq_mhz):
    # Find peaks in the data using a peak detector (good for more that 1 peak, min_dist needs work)
    indexes = peakutils.peak.indexes(spectrum.real,thres=0.2, min_dist=50)
    #indexes = signal.find_peaks_cwt(spectrum, numpy.arange(1,10))
             
    # Make an array of peak amplitudes and corresponding frequencies from detected peak indexes
    x = 0
    peaks = []
    freq = []
    while x < len(indexes):
        peaks.append(spectrum.real[indexes[x]])
        freq.append(faxis[indexes[x]])
        x += 1
        
    # search for peak at frequency closest to frequency of interest
    # subtract frequency of interest from frequencies the peaks are detected at.
    search = []
    for i in freq:
        search.append(i - freq_mhz)
    
    # convert values to absolute
    abs_search = numpy.abs(search)
    
    # lowest number will be closest to frequency of interest
    # find the index of this number in the list
    min_index = numpy.argmin(abs_search)
    
    # check peak is close enough to the frequency of interest
    if (freq_mhz-0.05) < freq[min_index] < (freq_mhz+0.05):
        
        
        amp_peak_1 = peaks[min_index]
        freq_peak_1 = freq[min_index]
            
        #print('Primary measurement used')            
        #print('Close Peak Detected\n')
        
    # if peak detected isn't close enough use bin average
    else:
        
        # search for bins closest to frequency of interest and save
        search_freq = []
        for i in faxis:
            search_freq.append(i - freq_mhz)
    
        # convert values to absolute
        abs_search_freq = numpy.abs(search_freq)
    
        # lowest number will be closest to frequency of interest
        # find the index of this number in the list
        min_index_freq = numpy.argmin(abs_search_freq)
                
        amp_peak_1 = spectrum[min_index_freq]
        freq_peak_1 = faxis[min_index_freq]
        
        #print('Central Bin Measurement Recorded') 
        #print('No Close Peak Detected\n')
    
    
    # take an average of the bins excluding peaks to estimate noise

    avg_bin = 0
    
    for i in spectrum:
        avg_bin = avg_bin + i
        
    for i in peaks:
        avg_bin = avg_bin - i 
  
    avg_bin = avg_bin/(len(spectrum)-len(peaks))
    
    #print('Primary Measurement:')    
    #print('%.3f dBm' % (peaks[min_index]-144))
    #print('%.4f MHz'% freq[min_index]) 
    
    #print('\nEstimated Noise Measurement:')    
    #print('%.3f dBm' % (avg_bin-144))
    #print('%.4f MHz'% faxis[2500]) 
        
    # order detected peaks from small to large
    peaks_index = numpy.argsort(peaks)
    
    # use peak index to retrieve data from data sets, insert dummy data if the peak index is not long enough (less than 4 peaks detected)
    x = 0
    
    if x < len(peaks_index):
        amp_peak_2 = peaks[peaks_index[len(peaks_index)-1]]
        freq_peak_2 = freq[peaks_index[len(peaks_index)-1]]
        
        x+=1
    
    else:
        amp_peak_2 = 0
        freq_peak_2 = 0
        
    if x < len(peaks_index):
        amp_peak_3 = peaks[peaks_index[len(peaks_index)-2]]
        freq_peak_3 = freq[peaks_index[len(peaks_index)-2]]
        
        x+=1
    
    else:
        amp_peak_3 = 0
        freq_peak_3 = 0
        
    if x < len(peaks_index):
        amp_peak_4 = peaks[peaks_index[len(peaks_index)-3]]
        freq_peak_4 = freq[peaks_index[len(peaks_index)-3]]
        
        x+=1
    
    else:
        amp_peak_4 = 0
        freq_peak_4 = 0
        
    
    # Save noise measurement 
    amp_peak_5 = avg_bin
    freq_peak_5 = 'noise'
        
    # use peak index to retrieve data from data sets
    #amp_peak_2 = peaks[peaks_index[len(peaks_index)-1]]
    #freq_peak_2 = freq[peaks_index[len(peaks_index)-1]]
    #amp_peak_3 = peaks[peaks_index[len(peaks_index)-2]]
    #freq_peak_3 = freq[peaks_index[len(peaks_index)-2]]
    #amp_peak_4 = peaks[peaks_index[len(peaks_index)-3]]
    #freq_peak_4 = freq[peaks_index[len(peaks_index)-3]]
    #amp_peak_5 = peaks[peaks_index[len(peaks_index)-4]]
    #freq_peak_5 = freq[peaks_index[len(peaks_index)-4]]

    return amp_peak_1, freq_peak_1, amp_peak_2, freq_peak_2, amp_peak_3, freq_peak_3, amp_peak_4, freq_peak_4, amp_peak_5, freq_peak_5, indexes

# Function to display results and write to a .CSV and .TXT file
def display_write(current_date, current_hour, amp_peak_1, freq_peak_1, amp_peak_2, freq_peak_2, amp_peak_3, freq_peak_3, amp_peak_4, freq_peak_4, amp_peak_5, freq_peak_5, lon, lat, alt):
    # print first 2 detected peaks
    print('%s %s Peaks: %.2f dB at %.3f MHz, %.2f dB at %.3f MHz, %.2f dB at %.3f MHz, %.2f dB at %.3f MHz and %.2f dB for %s \nLocation: Longitude %.3f, Latitude %.3f, Altitude %.3fm\n' % (current_date, current_hour, amp_peak_1, freq_peak_1, amp_peak_2, freq_peak_2, amp_peak_3, freq_peak_3, amp_peak_4, freq_peak_4, amp_peak_5, freq_peak_5, lon, lat, alt))

    # write peaks to a text file
    with open("log.txt","a") as f:
        f.write('%s %s Peaks: %.2f dB at %.3f MHz, %.2f dB at %.3f MHz, %.2f dB at %.3f MHz, %.2f dB at %.3f MHz and %.2f dB for %s Location: Longitude %.3f, Latitude %.3f, Altitude %.3fm\n' % (current_date, current_hour, amp_peak_1, freq_peak_1, amp_peak_2, freq_peak_2, amp_peak_3, freq_peak_3, amp_peak_4, freq_peak_4, amp_peak_5, freq_peak_5, lon, lat, alt))
        
    f.close()

    #write peaks to a .CSV file
    with open('log.csv', 'a') as csvfile:
        writer = csv.writer(csvfile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL, lineterminator = '\n')
        writer.writerow([current_date] + [current_hour] + [amp_peak_1] + [freq_peak_1] + [amp_peak_2] + [freq_peak_2] + [amp_peak_3] + [freq_peak_3] + [amp_peak_4] + [freq_peak_4] + [amp_peak_5] + [freq_peak_5] + [lon] + [lat] + [alt])

    csvfile.close()

    return

def fft(sdr, num_samples):
    
    try:
        # Read in samples from the device
        samples = sdr.read_samples(num_samples)
    
        w = numpy.hamming(num_samples)            
        # Process fft
        # take fft of samples
        fft = numpy.fft.fft(samples*w)
        
        # convert to dB
        spectrum = 20 * numpy.log10(abs(fft*2)) 
    
    except:
        spectrum = 0
                
    return spectrum

# Function to get samples and create FFT
def data(sdr, num_avg, num_samples):

    print("Data Capture in Progress........\n") 
    
    x = 1
    a = 1
    spectrum_avg = numpy.zeros(num_samples)
    
    while x <= num_avg:
        #time.sleep(0.1)
        spectrumA = fft(sdr, num_samples)
        spectrum_avg = spectrum_avg + spectrumA
        #print (spectrum_avg)
        #print (spectrumA)
        if a == 100:
            percent_comp = (x/num_avg)*100
            print("%d Averages Taken" % x)
            print("Sampling %.2f %% Complete\n"  % percent_comp)
            
            a = 1
        else:
            a += 1
        x += 1
     
    
    spectrum = spectrum_avg / num_avg
    #print (spectrum)
    
    #spectrum, faxis, samples = data(sdr)
    
    print("Data Capture Complete\n")    

    #print ("FFT Ready\n")
    
    # shift the fft so Fc is at the centre of the plot
    spectrum = numpy.fft.fftshift(spectrum)
                    
    array_length = len(spectrum)
        
    # create frequency axis, noting that Fc is in centre
    fstep = sdr.sample_rate/len(spectrum) # fft bin size
    fstep = fstep / 1000000 #  convert to MHz
        
    faxis = []
    k = 0
    while k < array_length:
        # create f axis with Fc in the centre
        faxis.append(((sdr.center_freq/1000000)-(fstep*(len(spectrum)/2))) + k * fstep)
        # below creates f axis with Fc at the start
        # faxis.append((sdr.center_freq/1000000) + k * fstep)
        k = k + 1


    return spectrum, faxis,spectrumA

# Function to set up and configure SDR
def sdr_control(freq,i,flag):
    
    try:
        
        # Create device
        sdr = RtlSdr()

        # configure device
        sdr.sample_rate = 2.048e6  # Hz
        sdr.center_freq = freq+0.1e6   # Hz
        sdr.gain = 49#25.4 #49 #0
        
        return sdr, flag
    
    except:
        
        print("RTL SDR Communication Error")
        
        flag = True
        
        if i <= 60:
            
            time.sleep(1)
            print("Attempting to Re-establish Connection...... Attempt: %d\n" %i)
            i = i + 1
            sdr, flag = sdr_control(freq,i,flag)
            
        else:
            
            # gpio disconnect
            try:
                gpio.cleanup()
            except:
                print("GPIO Not Deactivated")
            
            # exit programme with an error message
            sys.exit("RTL SDR Device Not Connected\n")
       
    return sdr, flag

# Function to set up Rpi gpio for antenna switching
def gpio_set():
    
    try:
        # pin numbering mode
        gpio.setmode(gpio.BOARD)

        # pins to be used
        chan_list = [16,18,22,24,26]

        # set pins as outputs with an initial value of low
        gpio.setup(chan_list, gpio.OUT, initial=gpio.LOW)
    
    except:
        print("GPIO Not Setup")

    return

# function to switch antenna
def gpio_switch(chan, state):
    
    try:

        # activate selected gpio
        gpio.output(chan, state)

        # wait for switch
        time.sleep(0.5)
    
    except:
        print("GPIO Not Switched")
    
    return

# Function to plot the FFT's to aid debugging
def debug_graph(faxis,spectrum,indexes, amp_peak_1, freq_peak_1, amp_peak_5, faxis2,spectrum2,indexes2, amp_peak_12, freq_peak_12, amp_peak_52):
    
    cal_factor_868 = -98 #-84 #-144
    cal_factor = -89
    
    try:
        
        freq_offset = 0.1
        peak_search_range = 0.05
        
        # plot the data from the 70MHz sample period
        plt.plot(faxis,(spectrum.real+cal_factor), label='Received Power')
        
        # plot noise level
        plt.plot([faxis[0],faxis[4999]],[(amp_peak_5+cal_factor),(amp_peak_5+cal_factor)], 'r--',linewidth=5, label='Estimated Noise Level')
        
        # determine min and max of y axis
        axes = plt.gca()
        ymin, ymax = axes.get_ylim()
        
        #Plot limits that peaks are searched for within
        plt.plot([(faxis[2500]-peak_search_range-freq_offset),(faxis[2500]-peak_search_range-freq_offset)],[ymax,ymin], 'k:',linewidth=2.5, label = 'Peak Search Limits')
        plt.plot([(faxis[2500]+peak_search_range-freq_offset),(faxis[2500]+peak_search_range-freq_offset)],[ymax,ymin], 'k:',linewidth=2.5)

        # Plot all detected peaks
        plot_peak_freq = []
        plot_peak_amp = []
        for i in indexes:
            plot_peak_freq.append(faxis[i])
            plot_peak_amp.append((spectrum.real[i]+cal_factor))
            
            
        plt.plot(plot_peak_freq, plot_peak_amp, 'go', label = 'Detected Peaks')
            
        # annotate plot with highest peak    
        plt.annotate('Observed Peak Power \n %.1f dBm \n %.3f MHz'%((amp_peak_1+cal_factor), freq_peak_1),
             xy=(freq_peak_1, (amp_peak_1+cal_factor)), xycoords='data',
             xytext=(0.65, 0.75), textcoords='figure fraction' ,
             arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=.2"))
        
        # annotate plot with noise value    
        plt.annotate('Estimated Noise \n %.1f dBm'%((amp_peak_5+cal_factor)),
             xy=(faxis[2500], (amp_peak_5+cal_factor)), xycoords='data',
             xytext=(0.2, 0.6), textcoords='figure fraction')

        # format plot
        plt.xlabel('Frequency (MHz)')
        plt.ylabel('Received Power (dBm)')
        plt.title('FFT of Spectrum')
        plt.grid()
        plt.legend(loc='upper left',fontsize ='small')
        plt.autoscale(enable=True, axis='x', tight=True)

        plt.figure()
        # plot the data from the last sample period
        plt.plot(faxis2,(spectrum2.real+cal_factor_868), label = 'Received Power')
        
        # plot noise level
        plt.plot([faxis2[0],faxis2[4999]],[(amp_peak_52+cal_factor_868),(amp_peak_52+cal_factor_868)], 'r--',linewidth=5, label = 'Estimated Noise Level')
        
        # determine min and max of y axis
        axes = plt.gca()
        ymin, ymax = axes.get_ylim()
        
        #Plot limits that peaks are searched for within
        plt.plot([(faxis2[2500]-peak_search_range-freq_offset),(faxis2[2500]-peak_search_range-freq_offset)],[ymax,ymin], 'k:',linewidth=2.5, label = 'Peak search Limits')
        plt.plot([(faxis2[2500]+peak_search_range-freq_offset),(faxis2[2500]+peak_search_range-freq_offset)],[ymax,ymin], 'k:',linewidth=2.5)

        # Plot all detected peaks
        plot_peak_freq2 = []
        plot_peak_amp2 = []
        for i in indexes2:
            plot_peak_freq2.append(faxis2[i])
            plot_peak_amp2.append((spectrum2.real[i]+cal_factor_868))
            
        plt.plot(plot_peak_freq2, plot_peak_amp2, 'go', label = 'Detected Peaks')

        # annotate plot with highest peak
        plt.annotate('Observed Peak Power \n %.1f dBm \n %.3f MHz'%((amp_peak_12+cal_factor_868), freq_peak_12),
             xy=(freq_peak_12, (amp_peak_12+cal_factor_868)), xycoords='data',
             xytext=(0.65, 0.75), textcoords='figure fraction',
             arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=.2"))
        
        # annotate plot with noise value    
        plt.annotate('Estimated Noise \n %.1f dBm'%((amp_peak_52+cal_factor_868)),
             xy=(faxis2[2500], (amp_peak_52+cal_factor_868)), xycoords='data',
             xytext=(0.2, 0.6), textcoords='figure fraction')

        # format plot
        plt.xlabel('Frequency (MHz)')
        plt.ylabel('Received Power (dBm)')
        plt.title('FFT of Spectrum')
        plt.grid()
        plt.legend(loc='upper left',fontsize ='small')
        plt.autoscale(enable=True, axis='x', tight=True)
        
        
        plt.show()
        
    except:
        
        print("\nComplete\n")
    
    return

# Function to save the unprocessed sampled data for possible later use
def raw_save(time_for_save,freq_mhz,sample_rate,samples):
    
    # convert sample rate to MHz
    sample_rate = sample_rate/1000000
    
    # save samples to .NPY file
    numpy.save('%s-%.1fMHz-%.4fMHz-raw_samples' % (time_for_save,freq_mhz,sample_rate) , samples)
    
    # print confirmation of file name used to terminal
    print ('%s-%.1fMHz-%.4fMHz-raw_samples file saved \n' % (time_for_save,freq_mhz,sample_rate))
    
    return 


## Main Body ##

# Get Location
def test(flag):
    
    #start_time = time.time()
    #print("--- %s seconds ---" % (time.time() - start_time))
    
    # Set up GPIO
    gpio_set()
    
    # Set first frequency of interest and display in MHz
    freq = 71e6
    freq_mhz = freq/1000000
    num_avg = 100
    num_samples = 5000
    print("Sampling at %.2f MHz\n" % freq_mhz)
    
    # Activate 70MHz antenna
    gpio_switch(18,1)
    
    #RTL re-connections counter
    i = 1
    
    # Set up SDR
    sdr, flag = sdr_control(freq,i,flag)
    
    # Find current GPS Location
    lon, lat, alt = gps_loc()
    
    # Collect Data 
    spectrum, faxis,samples = data(sdr, num_avg, num_samples)
    
    # Detect Peaks   
    amp_peak_1, freq_peak_1, amp_peak_2, freq_peak_2, amp_peak_3, freq_peak_3, amp_peak_4, freq_peak_4, amp_peak_5, freq_peak_5, indexes = peak_det(spectrum, faxis, freq_mhz)
    
    # Get current time and date    
    current_date, current_hour, current_time, time_for_save = time_date()
    
    # Save the sampled data for later use if needed 
    raw_save(time_for_save,freq_mhz,sdr.sample_rate, spectrum)
    
    # Record results
    display_write(current_date, current_hour, amp_peak_1, freq_peak_1, amp_peak_2, freq_peak_2, amp_peak_3, freq_peak_3, amp_peak_4, freq_peak_4, amp_peak_5, freq_peak_5, lon, lat, alt)
    
    try:
        # Disconnect SDR
        sdr.close()
        
    except:
        
        print('\nFailed to Close RTL SDR\n' )
    
    # Deactivate 70MHz antenna
    gpio_switch(18,0)
        
    # Taking second sample
    
    # Set second frequency of interest and display in MHz
    freq = 869.525e6
    freq_mhz = freq/1000000
    num_avg = 100
    num_samples = 5000
    print("Sampling at %.3f MHz\n" % freq_mhz)
    
    # Activate 868MHz antenna
    gpio_switch(16,1)
    
    #RTL re-connections counter
    i = 1
    
    # Setup SDR
    sdr, flag = sdr_control(freq,i,flag)
    
    # Find current GPS Location
    lon, lat, alt = gps_loc()
    
    # Collect Data 
    spectrum2, faxis2,samples2 = data(sdr, num_avg, num_samples)
    
    # Detect Peaks   
    amp_peak_12, freq_peak_12, amp_peak_22, freq_peak_22, amp_peak_32, freq_peak_32, amp_peak_42, freq_peak_42, amp_peak_52, freq_peak_52, indexes2 = peak_det(spectrum2, faxis2, freq_mhz)
    
    # Get current time and date     
    current_date, current_hour, current_time, time_for_save = time_date()
    
    # Save the sampled data for later use if needed 
    raw_save(time_for_save,freq_mhz,sdr.sample_rate,spectrum2)
    
    # Record results
    display_write(current_date, current_hour, amp_peak_12, freq_peak_12, amp_peak_22, freq_peak_22, amp_peak_32, freq_peak_32, amp_peak_42, freq_peak_42, amp_peak_52, freq_peak_52, lon, lat, alt)
    
    try:
        # Disconnect SDR
        sdr.close()
        
    except:
        
        print('\nFailed to Close RTL SDR\n' )
        
    
    # Deactivate 868MHz antenna
    gpio_switch(16,0)
    
    # Flash status lights to indicate successful completion
    gpio_switch([22,24,26],1)
    gpio_switch([22,24,26],0)
    
    # GPIO disconnect
    try:
        gpio.cleanup()
    except:
        print("GPIO Not Deactivated")
        
    # Plot FFT's    
    debug_graph(faxis,spectrum,indexes, amp_peak_1, freq_peak_1,amp_peak_5,faxis2,spectrum2,indexes2, amp_peak_12, freq_peak_12,amp_peak_52)
    
    return flag

#THIS PART ACTUALLY MAKES IT RUN#

# Schedule test to run every X minutes
#schedule.every(0.1).minutes.do(test)

# Attempt to connect to GPS device
try:
    # Connect to the local gpsd
    gpsd.connect()

# If device is unavailable, alert user.
except:
    print('GPS Not Available on This Device')

flag = False

# set up never ending loop
while True:
    # Check for pending tasks
    #schedule.run_pending()
    # Run test programme
    
    flag = test(flag)
        
   
    if flag == True:
        print('\nThis Device Has Recovered From a Communication Error\n' )
    
    if flag == False:
        print('\nThis Device Has Not Recovered From a Communication Error\n' )
    
    # Wait X seconds before next loop
    time.sleep(5)
    