#!/usr/bin/python

import os, sys, time, threading, datetime, math, csv
from PyQt4 import QtCore,QtGui,QtSql
from PyQt4.Qwt5 import Qwt
from gnuradio import gr, gru, optfir, eng_notation, blks2
import gnuradio.gr.gr_threading as _threading
import air_modes
from air_modes.modes_exceptions import *
from air_modes.modes_rx_ui import Ui_MainWindow
import sqlite3

class mainwindow(QtGui.QMainWindow):
    def __init__(self):
        QtGui.QMainWindow.__init__(self)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        #set defaults
        #add file, RTL, UHD sources
        self.ui.combo_source.addItems(["UHD device", "RTL-SDR", "File"])
        self.ui.combo_source.setCurrentIndex(0)

        #populate antenna, rate combo boxes based on source
        self.populate_source_options()

        #should round to actual achieved gain
        self.ui.line_gain.insert("30")

        #default to 3dB
        self.ui.line_threshold.insert("3")

        self.ui.prog_rssi.setMinimum(-40)
        self.ui.prog_rssi.setMaximum(0)

        self.ui.combo_ant.setCurrentIndex(self.ui.combo_ant.findText("RX2"))

        #check KML by default, leave the rest unchecked.
        self.ui.check_sbs1.setCheckState(QtCore.Qt.Unchecked)
        self.ui.check_raw.setCheckState(QtCore.Qt.Unchecked)
        self.ui.check_fgfs.setCheckState(QtCore.Qt.Unchecked)
        self.ui.check_kml.setCheckState(QtCore.Qt.Checked)

        self.ui.line_sbs1port.insert("30003")
        self.ui.line_rawport.insert("9988")
        self.ui.line_fgfsport.insert("5500")
        self.ui.line_kmlfilename.insert("modes.kml")

        #disable by default
        self.ui.check_adsbonly.setCheckState(QtCore.Qt.Unchecked)

        self.queue = gr.msg_queue(10)
        self.runner = None
        self.fg = None
        self.outputs = []
        self.updates = []
        self.output_handler = None
        self.kmlgen = None #necessary bc we stop its thread in shutdown
        self.dbinput = None
        self.dbname = "air_modes.db"

        self.datamodel = dashboard_data_model(None)
        self.ui.list_aircraft.setModel(self.datamodel)
        self.ui.list_aircraft.setModelColumn(0)

        #set up dashboard views
        #TODO: figure out why you can't update the RSSI or compass widgets with this mapper
        #pretty sure it's a Qt bug in the SQL code -- returns as string instead of double
        self.icaodelegate = ICAOViewDelegate()
        self.ui.list_aircraft.setItemDelegate(self.icaodelegate)
        self.dashboard_mapper = QtGui.QDataWidgetMapper()
        self.dashboard_mapper.setModel(self.datamodel)
        self.dashboard_mapper.addMapping(self.ui.line_icao, 0)
        #self.dashboard_mapper.addMapping(self.ui.prog_rssi, 2)
        self.dashboard_mapper.addMapping(self.ui.line_latitude, 3)
        self.dashboard_mapper.addMapping(self.ui.line_longitude, 4)
        self.dashboard_mapper.addMapping(self.ui.line_alt, 5)
        self.dashboard_mapper.addMapping(self.ui.line_speed, 6)
        #self.dashboard_mapper.addMapping(self.ui.compass_heading, 7)
        self.dashboard_mapper.addMapping(self.ui.line_climb, 8)
        self.dashboard_mapper.addMapping(self.ui.line_ident, 9)
        self.dashboard_mapper.addMapping(self.ui.line_type, 10)
        self.dashboard_mapper.addMapping(self.ui.line_range, 11)

        compass_palette = QtGui.QPalette()
        compass_palette.setColor(QtGui.QPalette.Foreground, QtCore.Qt.white)
        self.ui.compass_heading.setPalette(compass_palette)
        self.ui.compass_bearing.setPalette(compass_palette)
        #TODO: change the needle to an aircraft silhouette
        self.ui.compass_heading.setNeedle(Qwt.QwtDialSimpleNeedle(Qwt.QwtDialSimpleNeedle.Ray, False, QtCore.Qt.black))
        self.ui.compass_bearing.setNeedle(Qwt.QwtDialSimpleNeedle(Qwt.QwtDialSimpleNeedle.Ray, False, QtCore.Qt.black))
        #TODO FIXME do not release
        self.ui.line_my_lat.insert("37.7623")
        self.ui.line_my_lon.insert("-122.442527")

        #hook up the update signal
        self.ui.list_aircraft.selectionModel().currentRowChanged.connect(self.dashboard_mapper.setCurrentModelIndex)
        self.ui.list_aircraft.selectionModel().currentRowChanged.connect(self.update_heading_widget)
        self.ui.list_aircraft.selectionModel().currentRowChanged.connect(self.update_bearing_widget)
        self.datamodel.dataChanged.connect(self.compass_widgets_dataChanged)
        self.ui.list_aircraft.selectionModel().currentRowChanged.connect(self.update_rssi_widget)

    def update_heading_widget(self, index):
        if index.model() is not None:
            heading = index.model().data(index.model().index(index.row(), self.datamodel._colnames.index("heading"))).toDouble()[0]
            self.ui.compass_heading.setValue(heading)

    def update_bearing_widget(self, index):
        if index.model() is not None:
            bearing = index.model().data(index.model().index(index.row(), self.datamodel._colnames.index("bearing"))).toDouble()[0]
            self.ui.compass_bearing.setValue(bearing)

    def compass_widgets_dataChanged(self, startIndex, endIndex):
        index = self.ui.list_aircraft.selectionModel().currentIndex()
        if index.row() in range(startIndex.row(), endIndex.row()+1): #the current aircraft was affected
            if self.datamodel._colnames.index("heading") in range(startIndex.column(), endIndex.column()+1):
                self.update_heading_widget(index)
            if self.datamodel._colnames.index("bearing") in range(startIndex.column(), endIndex.column()+1):
                self.update_bearing_widget(index)

    def update_rssi_widget(self, index):
        if index.model() is not None:
            rssi = index.model().data(index.model().index(index.row(), 2)).toDouble()[0]
            self.ui.prog_rssi.setValue(rssi)

    #goes and gets valid antenna, sample rate options from the device and grays out appropriate things
    def populate_source_options(self):
        sourceid = self.ui.combo_source.currentText()
        self.rates = []
        self.ratetext = []
        self.antennas = []
        
        if sourceid == "UHD device":
            try:
                from gnuradio import uhd
                self.src = uhd.single_usrp_source("", uhd.io_type_t.COMPLEX_FLOAT32, 1)
                self.rates = [rate.start() for rate in self.src.get_samp_rates()]
                self.antennas = self.src.get_antennas()
                self.src = None #deconstruct UHD source for now
                self.ui.combo_ant.setEnabled(True)
                self.ui.combo_rate.setEnabled(True)
                self.ui.stack_source.setCurrentIndex(0)
            except:
                self.rates = []
                self.antennas = []
                self.ui.combo_ant.setEnabled(False)
                self.ui.combo_rate.setEnabled(False)
                self.ui.stack_source.setCurrentIndex(0)
                
        elif sourceid == "RTL-SDR":
            self.rates = [2.4e6]
            self.antennas = ["RX"]
            self.ui.combo_ant.setEnabled(False)
            self.ui.combo_rate.setEnabled(False)
            self.ui.stack_source.setCurrentIndex(0)

        elif sourceid == "File":
            self.rates = [2e6, 4e6, 6e6, 8e6, 10e6]
            self.antennas = ["None"]
            self.ui.combo_ant.setEnabled(False)
            self.ui.combo_rate.setEnabled(True)
            self.ui.stack_source.setCurrentIndex(1)

        self.ui.combo_rate.clear()
        self.ratetext = ["%.3f" % (rate / 1.e6) for rate in self.rates]
        for rate, text in zip(self.rates, self.ratetext):
            self.ui.combo_rate.addItem(text, rate)

        self.ui.combo_ant.clear()
        self.ui.combo_ant.addItems(self.antennas)

        if 4e6 in self.rates:
            self.ui.combo_rate.setCurrentIndex(self.rates.index(4e6))

    def on_combo_source_currentIndexChanged(self, index):
        self.populate_source_options()

    def on_button_start_released(self):
        #if we're already running, kill it!
        if self.runner is not None:
            self.output_handler.done = True
            self.output_handler = None
            self.outputs = []
            self.updates = []
            self.fg.stop()
            self.runner = None
            self.fg = None
            if self.kmlgen is not None:
                self.kmlgen.done = True
                #TODO FIXME need a way to kill kmlgen safely without delay
                #self.kmlgen.join()
                #self.kmlgen = None

            if self.dbinput is not None: self.dbinput.close()
            self.dbinput = None

            self.ui.button_start.setText("Start")

        else: #we aren't already running, let's get this party started
            options = {}
            options["source"] = str(self.ui.combo_source.currentText())
            options["rate"] = float(self.ui.combo_rate.currentText()) * 1e6
            options["antenna"] = str(self.ui.combo_ant.currentText())
            options["gain"] = float(self.ui.line_gain.text())
            options["threshold"] = float(self.ui.line_threshold.text())
            options["filename"] = str(self.ui.line_inputfile.text())

            self.fg = adsb_rx_block(options, self.queue) #create top RX block
            self.runner = top_block_runner(self.fg) #spawn new thread to do RX

            try:
                my_position = [float(self.ui.line_my_lat.text()), float(self.ui.line_my_lon.text())]
            except:
                my_position = None

            self.datamodelout = dashboard_output(my_position, self.datamodel)

            self.outputs = [self.datamodelout.output]
            self.updates = []

            #output options to populate outputs, updates
            if self.ui.check_kml.checkState():
                #we spawn a thread to run every 30 seconds (or whatever) to generate KML
                self.kmlgen = air_modes.modes_kml(self.ui.line_kmlfilename.text(), self.dbname, my_position) #create a KML generating thread

            if self.ui.check_sbs1.checkState():
                sbs1port = int(self.ui.line_sbs1port.text())
                sbs1out = air_modes.modes_output_sbs1(my_position, sbs1port)
                self.outputs.append(sbs1out.output)
                self.updates.append(sbs1out.add_pending_conns)

            if self.ui.check_fgfs.checkState():
                fghost = "127.0.0.1" #TODO FIXME
                fgport = self.ui.line_fgfsport.currentText()
                fgout = air_modes.modes_flightgear(my_position, fghost, int(fgport))
                self.outputs.append(fgout.output)

            if self.ui.check_raw.checkState():
                rawport = air_modes.modes_raw_server(int(self.ui.line_raw.text()))
                self.outputs.append(rawport.output)
                self.updates.append(rawport.add_pending_conns)

            self.livedata = air_modes.modes_output_print(my_position)
            #add output for live data box
            #TODO: this doesn't handle large volumes of data well; i get segfaults.
            #self.outputs.append(self.output_live_data)

            #create SQL database for KML and dashboard displays
            self.dbwriter = air_modes.modes_output_sql(my_position, self.dbname)
            self.outputs.append(self.dbwriter.output) #now the db will update itself

            #create output handler thread
            self.output_handler = output_handler(self.outputs, self.updates, self.queue)
            
            self.ui.button_start.setText("Stop")

    def output_live_data(self, msg):
        msgstr = self.livedata.parse(msg)
        if msgstr is not None:
            self.ui.text_livedata.append(msgstr)
            self.ui.text_livedata.verticalScrollBar().setSliderPosition(self.ui.text_livedata.verticalScrollBar().maximum())

#fades the ICAOs out as their last report gets older,
#and display ident if available, ICAO otherwise
class ICAOViewDelegate(QtGui.QStyledItemDelegate):
    def paint(self, painter, option, index):
        #draw selection rectangle
        if option.state & QtGui.QStyle.State_Selected:
            painter.setBrush(QtGui.QPalette().highlight())
            painter.drawRect(option.rect)

        #if there's an ident available, use it. otherwise print the ICAO
        if index.model().data(index.model().index(index.row(), 9)) != QtCore.QVariant():
            paintstr = index.model().data(index.model().index(index.row(), 9)).toString()
        else:
            paintstr = index.model().data(index.model().index(index.row(), 0)).toString()
        last_report = index.model().data(index.model().index(index.row(), 1)).toDouble()[0]
        age = (time.time() - last_report)
        max_age = 60. #age at which it grays out
        #minimum alpha is 0x40 (oldest), max is 0xFF (newest)
        age = min(age, max_age)
        alpha = int(0xFF - (0xBF / max_age) * age)
        painter.setPen(QtGui.QColor(0, 0, 0, alpha))
        painter.drawText(option.rect.left()+3, option.rect.top(), option.rect.width(), option.rect.height(), option.displayAlignment, paintstr)

#the data model used to display dashboard data.
class dashboard_data_model(QtCore.QAbstractTableModel):
    def __init__(self, parent):
        QtCore.QAbstractTableModel.__init__(self, parent)
        self._data = []
        self.lock = threading.Lock()
        self._colnames = ["icao", "seen", "rssi", "latitude", "longitude", "altitude", "speed", "heading", "vertical", "ident", "type", "range", "bearing"]
        #custom precision limits for display
        self._precisions = [None, None, None, 6, 6, 0, 0, 0, 0, None, None, 2, 0]
        for field in self._colnames:
            self.setHeaderData(self._colnames.index(field), QtCore.Qt.Horizontal, field)
    def rowCount(self, parent=QtCore.QVariant()):
        return len(self._data)
    def columnCount(self, parent=QtCore.QVariant()):
        return len(self._colnames)
    def data(self, index, role=QtCore.Qt.DisplayRole):
        if not index.isValid():
            return QtCore.QVariant()
        if index.row() >= self.rowCount():
            return QtCore.QVariant()
        if index.column() >= self.columnCount():
            return QtCore.QVariant()
        if (role != QtCore.Qt.DisplayRole) and (role != QtCore.Qt.EditRole):
            return QtCore.QVariant()
        if self._data[index.row()][index.column()] is None:
            return QtCore.QVariant()
        else:
            #if there's a dedicated precision for that column, print it out with the specified precision.
            #this only works well if you DON'T have other views/widgets that depend on numeric data coming out.
            #i don't like this, but it works for now. unfortunately it seems like Qt doesn't give you a
            #good alternative.
            if self._precisions[index.column()] is not None:
                return QtCore.QVariant("%.*f" % (self._precisions[index.column()], self._data[index.row()][index.column()]))
            else:
                if self._colnames[index.column()] == "icao":
                    return QtCore.QVariant("%06x" % self._data[index.row()][index.column()]) #return as hex string
                else:
                    return QtCore.QVariant(self._data[index.row()][index.column()])

    def setData(self, index, value, role=QtCore.Qt.EditRole):
        self.lock.acquire()
        if not index.isValid():
            return False
        if index.row() >= self.rowCount():
            return False
        if index.column >= self.columnCount():
            return False
        if role != QtCore.Qt.EditRole:
            return False
        self._data[index.row()][index.column()] = value
        self.lock.release()

    #addRecord implements an upsert on self._data; that is,
    #it updates the row if the ICAO exists, or else it creates a new row.
    def addRecord(self, record):
        self.lock.acquire()
        icaos = [x[0] for x in self._data]
        if record["icao"] in icaos:
            row = icaos.index(record["icao"])
            for column in record:
                self._data[row][self._colnames.index(column)] = record[column]
            #create index to existing row and tell the model everything's changed in this row
            #or inside the for loop, use dataChanged on each changed field (might be better)
            self.dataChanged.emit(self.createIndex(row, 0), self.createIndex(row, len(self._colnames)-1))

        #only create records for ICAOs with ADS-B reports
        elif ("latitude" or "speed" or "ident") in record:
            #find new inserted row number
            icaos.append(record["icao"])
            newrowoffset = sorted(icaos).index(record["icao"])
            self.beginInsertRows(QtCore.QModelIndex(), newrowoffset, newrowoffset)
            newrecord = [None for x in xrange(len(self._colnames))]
            for col in xrange(0, len(self._colnames)):
                if self._colnames[col] in record:
                    newrecord[col] = record[self._colnames[col]]
            self._data.append(newrecord)
            self._data = sorted(self._data, key = lambda x: x[0]) #sort by icao
            self.endInsertRows()
        self.lock.release()
        self.prune()

    #weeds out ICAOs older than 5 minutes
    def prune(self):
        self.lock.acquire()
        for (index,row) in enumerate(self._data):
            if time.time() - row[1] >= 60:
                self.beginRemoveRows(QtCore.QModelIndex(), index, index)
                self._data.pop(index)
                self.endRemoveRows()
        self.lock.release()
                

class dashboard_output(air_modes.modes_parse.modes_parse):
    def __init__(self, mypos, model):
        air_modes.modes_parse.modes_parse.__init__(self, mypos)
        self.model = model
    def output(self, msg):
        [data, ecc, reference, timestamp] = msg.split()
        data = air_modes.modes_parse.modes_reply(long(data, 16))
        ecc = long(ecc, 16)
        rssi = 10.*math.log10(float(reference))
        msgtype = data["df"]
        now = time.time()
        newrow = {"rssi": rssi, "seen": now}
        if msgtype in [0, 4, 20]:
            newrow["altitude"] = air_modes.altitude.decode_alt(data["ac"], True)
            newrow["icao"] = ecc
            self.model.addRecord(newrow)
        
        elif msgtype == 17:
            icao = data["aa"]
            newrow["icao"] = icao
            subtype = data["ftc"]
            if subtype == 4:
                (ident, actype) = self.parseBDS08(data)
                newrow["ident"] = ident
                newrow["type"] = actype
            elif 5 <= subtype <= 8:
                (ground_track, decoded_lat, decoded_lon, rnge, bearing) = self.parseBDS06(data)
                newrow["heading"] = ground_track
                newrow["latitude"] = decoded_lat
                newrow["longitude"] = decoded_lon
                newrow["altitude"] = 0
                if rnge is not None:
                    newrow["range"] = rnge
                    newrow["bearing"] = bearing
            elif 9 <= subtype <= 18:
                (altitude, decoded_lat, decoded_lon, rnge, bearing) = self.parseBDS05(data)
                newrow["altitude"] = altitude
                newrow["latitude"] = decoded_lat
                newrow["longitude"] = decoded_lon
                if rnge is not None:
                    newrow["range"] = rnge
                    newrow["bearing"] = bearing
            elif subtype == 19:
                subsubtype = data["sub"]
                velocity = None
                heading = None
                vert_spd = None
                if subsubtype == 0:
                    (velocity, heading, vert_spd) = self.parseBDS09_0(data)
                elif 1 <= subsubtype <= 2:
                    (velocity, heading, vert_spd) = self.parseBDS09_1(data)
                newrow["speed"] = velocity
                newrow["heading"] = heading
                newrow["vertical"] = vert_spd

            self.model.addRecord(newrow)


#the output handler is a thread which runs the various registered output functions when there's a received message.
#it also executes registered updates at the sleep rate -- currently 10Hz.
class output_handler(threading.Thread):
    def __init__(self, outputs, updates, queue):
        threading.Thread.__init__(self)
        self.setDaemon(1)
        self.outputs = outputs
        self.updates = updates
        self.queue = queue
        self.done = False
        self.start()

    def run(self):
        while self.done is False:
            for update in self.updates:
                update()
            while not self.queue.empty_p():
                msg = self.queue.delete_head()
                for output in self.outputs:
                    try:
                        output(msg.to_string())
                    except ADSBError:
                        pass

            time.sleep(0.1)

        self.done = True
        self.outputs = None
        self.updates = None
        self.queue = None
        

class top_block_runner(_threading.Thread):
    def __init__(self, tb):
        _threading.Thread.__init__(self)
        self.setDaemon(1)
        self.tb = tb
        self.done = False
        self.start()

    def run(self):
        self.tb.run()
        self.done = True

#Top block for ADSB receiver. If you define a standard interface you
#can make this common code between the GUI app and the cmdline app
class adsb_rx_block (gr.top_block):
    def __init__(self, options, queue):
        gr.top_block.__init__(self)

        self.options = options
        rate = options["rate"]
        print "Rate: %f" % rate
        use_resampler = False
        freq = 1090e6

        if options["source"] == "UHD device":
            from gnuradio import uhd
            self.u = uhd.single_usrp_source("", uhd.io_type_t.COMPLEX_FLOAT32, 1)
            time_spec = uhd.time_spec(0.0)
            self.u.set_time_now(time_spec)
            self.u.set_antenna(options["antenna"])
            self.u.set_samp_rate(rate)
            rate = self.u.get_samp_rate()
            self.u.set_gain(int(options["gain"]))
            self.u.set_center_freq(freq, 0)
        
        elif options["source"] == "RTL-SDR":
            import osmosdr
            self.u = osmosdr.source_c()
            self.u.set_sample_rate(2.4e6) #fixed for RTL dongles
            rate = int(4e6)
            self.u.set_gain_mode(0) #manual gain mode
            self.u.set_gain(int(options["gain"]))
            self.u.set_center_freq(freq, 0)
            use_resampler = True

        elif options["source"] == "File":
            self.u = gr.file_source(gr.sizeof_gr_complex, options["filename"])
        else:
            raise NotImplementedError

        self.demod = gr.complex_to_mag()
        self.avg = gr.moving_average_ff(100, 1.0/100, 400)

        self.preamble = air_modes.modes_preamble(int(rate), float(options["threshold"]))
        self.slicer = air_modes.modes_slicer(int(rate), queue)

        if use_resampler:
            self.lpfiltcoeffs = gr.firdes.low_pass(1, 5*2.4e6, 1.2e6, 300e3)
            self.resample = blks2.rational_resampler_ccf(interpolation=5, decimation=3, taps=self.lpfiltcoeffs)
            self.connect(self.u, self.resample, self.demod)
        else:
            self.connect(self.u, self.demod)

        self.connect(self.demod, self.avg)
        self.connect(self.demod, (self.preamble, 0))
        self.connect(self.avg, (self.preamble, 1))
        self.connect(self.preamble, self.slicer)
        

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    window = mainwindow()
    window.setWindowTitle("Mode S/ADS-B receiver")
    window.show()
    sys.exit(app.exec_())

