import sys
import numpy as np
from PyQt4 import QtGui, QtCore
from PyQt4.uic import loadUiType
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT as NavigationToolbar

from obspy import read
import matplotlib.pyplot as plt

Ui_MainWindow, QMainWindow = loadUiType('ProsSpect_winMain.ui')

listDataFileToOpen = []
fileName = []

# nD = [] # number of data in each sampling rate
timeLine = []
# year = []
# month = []
# day = []

flagLine = []
# hour = []
# minute = []
# second = []

Ex = []  # channel 1
Ey = []  # channel 2
Hx = []  # channel 3
Hy = []  # channel 4
Hz = []  # channel 5

class Main(QMainWindow, Ui_MainWindow):
    def __init__(self, ):
        super(Main, self).__init__()
        self.setupUi(self)
        self.setGeometry(500, 500, 500, 500)
        self.showMaximized()
        self.setWindowTitle('ProsPect')

        self.triggeredMenuAction()
        self.addFigure()
        self.tool()

    def triggeredMenuAction(self):
        self.actionImport_Signal.triggered.connect(self.importSignalBrowser)
        self.actionSave_Fi_gure.triggered.connect(self.saveFigure)

    def importSignalBrowser(self):
        nameFiletoOpen = QtGui.QFileDialog.getOpenFileNames(self, 'Import Signal', '', 'Phoenix ASCII Time Series (*.AS*)')
        if nameFiletoOpen:
            listDataFileToOpen.extend(nameFiletoOpen)
            self.addListWidgetListSignal()
            self.widgetListSignal.itemPressed.connect(self.item_Pressed)

    def tool(self):
        self.toolBar = self.addToolBar('Toolbar')
        self.toolBar.addAction(self.actionImport_Signal)
        self.toolBar = self.addToolBar('Figure')
        self.toolBar.addAction(self.actionSave_Fi_gure)

    def saveFigure(self):
        nameFiguretoSave = QtGui.QFileDialog.getSaveFileName(self, 'Save Figure', '', 'PNG files (*.png)')
        self.saveFig = plt.savefig(nameFiguretoSave)

    def addFigure(self):
        self.figure = plt.figure()
        # self.figure.patch.set_facecolor('#D6D6D6')
        self.canvas = FigureCanvas(self.figure)
        self.toolbar = NavigationToolbar(self.canvas, self)
        self.verticalLayout.addWidget(self.toolbar)
        self.verticalLayout.addWidget(self.canvas)
        self.canvas.draw()

    def addListWidgetListSignal(self):
        for i in range(len(listDataFileToOpen)):
            tmp = listDataFileToOpen[i].split('/')
            fileName.append(tmp[len(tmp) - 1])
            root = QtGui.QTreeWidgetItem(self.widgetListSignal, [str(fileName[i]),
                                                                 str(listDataFileToOpen[i])])

    def removeList(self):
        listDataFileToOpen.remove(self.nameToOpen)
        self.widgetListSignal.clear()
        self.addListWidgetListSignal()

    def item_Pressed(self, item):
        self.nameToOpen = str(item.text(1))
        self.ID = str(item.text(0))

    def contextMenuEvent(self, position):

        level = 0
        indexes = self.widgetListSignal.selectedIndexes()

        if len(indexes) > 0:
            level = 1

        if level == 1:
            menu = QtGui.QMenu(self)

            InsertSig = menu.addAction('Import Signal')
            InsertSig.triggered.connect(self.importSignalBrowser)
            RemoveSig = menu.addAction('Remove Signal')
            RemoveSig.triggered.connect(self.removeList)

            menu.addSeparator()

            TimeSeries = menu.addMenu('Time Series')
            TimeSeries1 = TimeSeries.addAction('Plot')
            TimeSeries1.triggered.connect(self.plotTimeSeries1)

            Spectogram = menu.addMenu('Spectogram')
            Spectogram1 = Spectogram.addAction('Plot')
            Spectogram1.triggered.connect(self.plotSpectogram1)


            menu.exec_(QtGui.QCursor.pos())

    def detrend(self, t, channel, order = 6): # ___________________________________________________________________________________
        # function detrend using polinomial
        # default orde = 6
        # (t,channel, order = 6)

        model = np.polyfit(t, channel, order)
        dtrnd = channel - np.polyval(model, t)
        return dtrnd

    def plotTimeSeries1(self):
        self.namePlotTo_1 = self.nameToOpen
        fileName, timeLine, flagLine, SR, Ex, Ey, Hx, Hy, Hz = self.openSignal(self.namePlotTo_1)

        dt = 1. / SR
        nData = len(Ex)
        t = np.arange(0.0, nData / SR, dt)

        print('Trend removal ...')
        Ex = self.detrend(t, Ex)
        print('- Ex')
        Ey = self.detrend(t, Ey)
        print('- Ey')
        Hx = self.detrend(t, Hx)
        print('- Hx')
        Hy = self.detrend(t, Hy)
        print('- Hy')
        Hz = self.detrend(t, Hz)
        print('- Hz')

        self.rmmpl()
        self.addFigure()

        pltsf1 = self.figure.add_subplot(511)
        pltsf1.hold(False)
        pltsf1.plot(t, Ex)
        pltsf1.grid(True)
        pltsf1.axis([min(t), max(t), -50000, 50000])
        pltsf1.set_ylabel('Ex')

        pltsf2 = self.figure.add_subplot(512)
        pltsf2.hold(False)
        pltsf2.plot(t, Ey)
        pltsf2.grid(True)
        pltsf2.axis([min(t), max(t), -50000, 50000])
        pltsf2.set_ylabel('Ey')

        pltsf3 = self.figure.add_subplot(513)
        pltsf3.hold(False)
        pltsf3.plot(t, Hx)
        pltsf3.grid(True)
        pltsf3.axis([min(t), max(t), -100000, 100000])
        pltsf3.set_ylabel('Hx')

        pltsf4 = self.figure.add_subplot(514)
        pltsf4.hold(False)
        pltsf4.plot(t, Hy)
        pltsf4.grid(True)
        pltsf4.axis([min(t), max(t), -100000, 100000])
        pltsf4.set_ylabel('Hy')

        pltsf5 = self.figure.add_subplot(515)
        pltsf5.hold(False)
        pltsf5.plot(t, Hz)
        pltsf5.grid(True)
        pltsf5.axis([min(t), max(t), -100000, 100000])
        pltsf5.set_ylabel('Hz')

        xticklabels = pltsf1.get_xticklabels() + pltsf2.get_xticklabels() + pltsf3.get_xticklabels() + pltsf4.get_xticklabels()
        plt.setp(xticklabels, visible=False)

        pltsf5.set_xlabel('Time [s]')

        self.figure.tight_layout()
        self.canvas.draw()

    def plotSpectogram1(self):

        self.namePlotTo_1 = self.nameToOpen
        fileName, timeLine, flagLine, SR, Ex, Ey, Hx, Hy, Hz = self.openSignal(self.namePlotTo_1)

        nData = len(Ex)

        xMin = 0
        xMax = nData / SR

        self.rmmpl()
        self.addFigure()

        sp1 = self.figure.add_subplot(321)  # create an axis
        sp1.hold(False)  # discards the old graph
        PEx, freqs, bins, im = sp1.specgram(Ex, NFFT=8192, Fs=SR, noverlap=4906)
        sp1.axis([xMin, xMax, min(freqs), max(freqs)])
        sp1.set_xlabel('t [s]')
        sp1.set_ylabel('Frequency [Hz]')
        sp1.set_title('Ex')

        sp2 = self.figure.add_subplot(323)  # create an axis
        sp2.hold(False)  # discards the old graph
        PEx, freqs, bins, im = sp2.specgram(Ey, NFFT=8192, Fs=SR, noverlap=4906)
        sp2.axis([xMin, xMax, min(freqs), max(freqs)])
        sp2.set_xlabel('t [s]')
        sp2.set_ylabel('Frequency [Hz]')
        sp2.set_title('Ey')

        sp3 = self.figure.add_subplot(322)  # create an axis
        sp3.hold(False)  # discards the old graph
        PEx, freqs, bins, im = sp3.specgram(Hx, NFFT=8192, Fs=SR, noverlap=4906)
        sp3.axis([xMin, xMax, min(freqs), max(freqs)])
        sp3.set_xlabel('t [s]')
        sp3.set_ylabel('Frequency [Hz]')
        sp3.set_title('Hx')

        sp4 = self.figure.add_subplot(324)  # create an axis
        sp4.hold(False)  # discards the old graph
        PEx, freqs, bins, im = sp4.specgram(Hy, NFFT=8192, Fs=SR, noverlap=4906)
        sp4.axis([xMin, xMax, min(freqs), max(freqs)])
        sp4.set_xlabel('t [s]')
        sp4.set_ylabel('Frequency [Hz]')
        sp4.set_title('Hy')

        sp5 = self.figure.add_subplot(326)  # create an axis
        sp5.hold(False)  # discards the old graph
        PEx, freqs, bins, im = sp5.specgram(Hz, NFFT=8192, Fs=SR, noverlap=4906)
        sp5.axis([xMin, xMax, min(freqs), max(freqs)])
        sp5.set_xlabel('t [s]')
        sp5.set_ylabel('Frequency [Hz]')
        sp5.set_title('Hz')

        # xticklabels = sp1.get_xticklabels() + sp2.get_xticklabels() + sp3.get_xticklabels()
        # plt.setp(xticklabels, visible=False)

        self.figure.tight_layout()
        self.canvas.draw()

    def rmmpl(self,):
        self.verticalLayout.removeWidget(self.canvas)
        self.canvas.close()
        self.verticalLayout.removeWidget(self.toolbar)
        self.toolbar.close()

    def openSignal(self, namePlot):

        print('Read data', namePlot, '...')
        # this function are used to read data phoenix time series acsii
        # the return value was SR, Ex, Ey, Hx, Hy, Hz

        openFile = open(namePlot, 'r')

        fileName = openFile.readline()

        while True:
            read = openFile.readline()

            if read == '':
                # either end of file or just a blank line.....
                # we'll assume EOF, because we don't have a choice with the while loop!
                break

            # read time
            timeSplit = read.split()
            timeLine.append(read)

            # get sampling rate
            tmp = timeSplit[6].split(':')
            SR = int(tmp[1])

            # # get date
            # tmp = timeSplit[2].split('/')
            # for i in range(0, SR):
            #     year.append(tmp[0])
            #     month.append(tmp[1])
            #     day.append(tmp[2])
            #
            # # get time
            # tmp = timeSplit[4].split(':')
            # for i in range(0, SR):
            #     hour.append(tmp[0])
            #     minute.append(tmp[1])
            #     second.append(tmp[2])

            # channel Flag
            read = openFile.readline()
            flagLine.append(read)

            # read the channel
            for i in range(0, SR):
                read = openFile.readline()
                tmp = read.split()
                # nD.append(int(tmp[0]))
                Ex.append(int(tmp[1]))
                Ey.append(int(tmp[2]))
                Hx.append(int(tmp[3]))
                Hy.append(int(tmp[4]))
                Hz.append(int(tmp[5]))

        openFile.close()

        print('Done!')
        # return year, month, day, hour, minute, second, SR, Ex, Ey, Hx, Hy, Hz

        openFile.close()
        return fileName, timeLine, flagLine, SR, Ex, Ey, Hx, Hy, Hz

# ---------------------------------------------------------------------------------------------------------------------

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    main = Main()
    main.show()
    sys.exit(app.exec_())
