# -*- coding: utf-8 -*-

import sys
from io import BytesIO
from PySide import QtGui,QtCore
import scipy as spy
from .EQSignal_UI_M_D import Ui_EQSignal_
from .About_D import Ui_About_
from .EQSignal import EQSignal
from copy import deepcopy


XS = {True:'log',False:'linear'}

class AutoColor(object):
    colors = ["b", "g", "r", "c", "m", "y"]
    n = 6
    i = -1
    color = "b"

    def __iter__(self):
        return self

    def next(self):
        self.i += 1
        if self.i > self.n-1:
            self.i = 0
        self.color = self.colors[self.i]
        return self.color

    def reset(self):
        self.i = -1

class Ui_About(QtGui.QWidget,Ui_About_):
    """docstring for Ui_ABOUT"""
    def __init__(self, parent=None):
        super(Ui_About, self).__init__(parent)
        self.setupUi(self)

class Ui_EQSignal(QtGui.QMainWindow, Ui_EQSignal_):

    """Ui_EQSignal"""

    def __init__(self, parent=None):
        super(Ui_EQSignal, self).__init__(parent)
        self.setupUi(self)
        self.eqs = EQSignal()
        self.initViewTH()
        self.initViewSPA()
        self.initViewSP()
        self.initViewFFT()
        # self.initTable()
        self.fileName = ""
        self.about = Ui_About()
        self.textEditSPA = QtGui.QTextBrowser()
        self.textEditSPA.resize(400,600)
        self.textEditSPA.setWindowTitle("Absolute Acceleration Spectrum")
        self.textEditTH = QtGui.QTextBrowser()
        self.textEditTH.resize(800,600)
        self.textEditTH.setWindowTitle("Time History")

        self.setupConnection()

    def setupConnection(self):

        self.actionNGA_Record.triggered.connect(self.readNGAthenPlotTH)
        self.actionPlain_Text.triggered.connect(self.readTXTthenPlotTH)
        self.actionExcel.triggered.connect(self.eqs.save_xlsx)
        self.actionAbout.triggered.connect(self.about.show)

        self.ApplyAdjust.clicked.connect(self.adjustthenPlotTH)
        self.RSPA.clicked.connect(self.calSPthenPlotSP)
        self.ViewASP.clicked.connect(self.calASPthenPlotASP)
        self.RawON.stateChanged.connect(self.plotTH)

        self.OPEN.clicked.connect(self.openFile)
        self.Reload.clicked.connect(self.reloadFile)
        self.Norm.clicked.connect(self.NormthenPlotTH)
        self.RecoverN.clicked.connect(self.recoverthenPlotTH)
        self.Refresh.clicked.connect(self.plotTH)

        self.RePlotSPA.clicked.connect(self.plotCurrentSPA)

        self.setTB.clicked.connect(self.getTrimIndice)
        self.DoTrim.clicked.connect(self.trimthenPlotTH)
        self.Recover.clicked.connect(self.recoverthenPlotTH)

        self.X_LIN_FFT.clicked.connect(self.changeXS_FFT)
        self.X_LOG_FFT.clicked.connect(self.changeXS_FFT)
        self.X_LIN_SPA.clicked.connect(self.changeXS_SPA)
        self.X_LOG_SPA.clicked.connect(self.changeXS_SPA)
        self.X_LIN_SP.clicked.connect(self.changeXS_SP)
        self.X_LOG_SP.clicked.connect(self.changeXS_SP)

        self.SPfit.clicked.connect(self.fitSpectra)
        self.VDTH.clicked.connect(self.viewDataTH)
        self.VDSP.clicked.connect(self.viewDataSP)
        self.CDR.activated.connect(self.viewDataSPA)
        self.CDR2.activated.connect(self.viewDataSP)
        self.FSP.activated.connect(self.FFTPSD)
        self.DR.textChanged.connect(self.updateCDR)

        self.SPATON.stateChanged.connect(self.plotSPA)

    def initViewTH(self):

        fig = self.ViewTH.getFigure()
        accPlot = fig.add_subplot(3,1,1)
        accPlot.set_title("Acceleration")
        accPlot.axhline(0.0,color="0.8")
        accPlot.grid(True)
        
        velPlot = fig.add_subplot(3,1,2)
        velPlot.set_title("Velocity")
        velPlot.axhline(0.0,color="0.8")
        velPlot.grid(True)
        
        dispPlot = fig.add_subplot(3,1,3)
        dispPlot.set_title("Displacement")
        dispPlot.axhline(0.0,color="0.8")
        dispPlot.grid(True)

        dispPlot.set_xlabel("time (s)")

    def initViewFFT(self):

        fig = self.ViewFFT.getFigure()
        fftPlot = fig.add_subplot(1,1,1)

        fftPlot.set_title("Fourier Amplitude Spectrum")
        fftPlot.grid(True)
        fftPlot.set_xlabel("Frequency (Hz)")
        fftPlot.set_ylabel("Fourier Amplitude")

    def initViewSPA(self):

        fig = self.ViewSPA.getFigure()
        spPlot = fig.add_subplot(1,1,1)

        spPlot.set_title("Acceleration Spectra")
        # spPlot.set_xscale("log")
        spPlot.grid(True)
        spPlot.set_xlabel("Period (s)")
        spPlot.set_ylabel("Response Acceleration")
            
    def initViewSP(self):

        fig = self.ViewSP.getFigure()

        titles = ["Acceleration","Velocity","Displacement","Input Energy"]

        for i in range(4):
            spPlot = fig.add_subplot(2,2,i+1)
            spPlot.set_title(" ".join([titles[i],"Spectra"]))
            # spPlot.set_xscale("log")
            spPlot.grid(True)

    def setDefaultValue(self):
        self.Freq1.setValue(round(1.0/(self.eqs.dt*self.eqs.n),3))
        self.Freq2.setValue(round(0.5/self.eqs.dt))
        self.TS.setValue(round(2.5*self.eqs.dt,3))
    
    def viewDataTH(self):

        data = spy.vstack((self.eqs.t,self.eqs.acc,self.eqs.vel,self.eqs.disp)).T
        Data = BytesIO()
        spy.savetxt(Data,data,fmt="%12.6f",
            header="%12s\t%12s\t%12s\t%12s"%("Time","Acc",
                "Vel","Disp"),
            delimiter="\t")
        self.textEditTH.clear()
        self.textEditTH.setPlainText(Data.getvalue().decode()[2:])
        Data.close()
        self.textEditTH.show()

    def viewDataSP(self):

        i = self.CDR2.currentIndex()
        data = spy.vstack((self.eqs.sp[i].SPT,self.eqs.sp[i].SPA,self.eqs.sp[i].SPV,self.eqs.sp[i].SPD,self.eqs.sp[i].SPE)).T
        Data = BytesIO()
        spy.savetxt(Data,data,fmt="%12.6f",
            header="%12s\t%12s\t%12s\t%12s\t%12s"%("Period","SPA","SPV","SPD","SPE"),
            delimiter="\t")
        self.textBrowserSP.clear()
        self.textBrowserSP.setPlainText(Data.getvalue().decode()[2:])
        self.tabWidget.setCurrentIndex(4)
        Data.close()

    def viewDataSPA(self):
        
        i = self.CDR.currentIndex()

        data = spy.vstack((self.eqs.sp[i].SPT,self.eqs.sp[i].SPA)).T
        Data = BytesIO()
        spy.savetxt(Data,data,fmt="%14.8f",
            header="%14s\t%14s"%("Period","SPA"),
            delimiter="\t")
        self.textEditSPA.clear()
        self.textEditSPA.setPlainText(Data.getvalue().decode()[2:])
        self.textEditSPA.show()

        self.textEditSPA.setWindowTitle("zeta=%.2f"%self.eqs.sp[i].Zeta)
        # self.tableSPA.show()
        Data.close()

    def NormthenPlotTH(self):
        self.eqs.norm()
        self.plotTH()
    
    def updateCDR(self):
        self.CDR.clear()
        self.CDR2.clear()
        self.CDRfit.clear()
        DRs = self.DR.text().replace(" ","").split(",")
        self.CDR.addItems(DRs)
        self.CDRfit.addItems(DRs)
        DRs = ["Damping Ratio : "+dr for dr in DRs]
        self.CDR2.addItems(DRs)
    
    def openFile(self):

        fd = QtGui.QFileDialog()
        self.fileName = fd.getOpenFileName(self, "Open File", "./data/", "*.at2 *.txt")[0]
        self.URL.setText(self.fileName)
        self.reloadFile()

    def reloadFile(self):

        self.fileName = self.URL.text()

        if self.fileName[-4:] == ".txt":
            dt = QtGui.QInputDialog.getDouble(self, "Time Interval", "dt = ", value=0.02, decimals=3)[0]
            self.eqs.readTXTFile(self.fileName,dt)
            self.plotTH()
            self.tabWidget.setCurrentIndex(0)
            self.VDTH.setEnabled(True)

        elif self.fileName[-4:] == ".at2":
            self.eqs.readNGAFile(self.fileName)
            self.plotTH()
            self.tabWidget.setCurrentIndex(0)
            self.VDTH.setEnabled(True)
        else:
            pass

        self.setDefaultValue()

    def readNGAthenPlotTH(self):
        self.getNGAFileName()
        self.eqs.readNGAFile(self.fileName)
        self.plotTH()
        self.tabWidget.setCurrentIndex(0)
        self.VDTH.setEnabled(True)
        self.Freq1.setValue(round(1.0/(self.eqs.dt*self.eqs.n),2))
        self.Freq2.setValue(round(0.5/self.eqs.dt))
        self.TS.setValue(round(2.0*self.eqs.dt,2))

    def readTXTthenPlotTH(self):
        self.getTXTFileName()

        dt = QtGui.QInputDialog.getDouble(self, "Time Interval", "dt = ", value=0.02, decimals=3)[0]
        self.eqs.readTXTFile(self.fileName,dt)
        self.plotTH()
        self.tabWidget.setCurrentIndex(0)
        self.VDTH.setEnabled(True)

        self.Freq1.setValue(round(1.0/(self.eqs.dt*self.eqs.n),2))
        self.Freq2.setValue(round(0.5/self.eqs.dt))
        self.TS.setValue(round(2.0*self.eqs.dt,2))

    def getTrimIndice(self):
        thd1 = self.THD1.value()
        thd2 = self.THD2.value()

        iStart, iStop = self.eqs.getTrimIndice(thd1, thd2)
        self.iStart.setValue(iStart)
        self.iStop.setValue(iStop)

    def trimthenPlotTH(self):
        iStart = max(self.iStart.value()-1,0)
        iStop  = min(self.iStop.value()-1,self.eqs.n-1)
        self.eqs.trim(iStart,iStop)
        self.plotTH()

    def recoverthenPlotTH(self):
        self.eqs.recover()
        self.plotTH()

    def adjustthenPlotTH(self):

        filt = self.Filt.isChecked()
        detr = self.Detrend.isChecked()
        tadj = self.TAdjust.isChecked()

        curr = False

        if filt:
            order = self.FOrder.value()
            f1 = self.Freq1.value()
            f2 = self.Freq2.value()
            btype = self.BandType.currentText()
            self.eqs.filt(order, f1, f2, btype, curr)
            curr = True

        if detr:
            oh = self.OHigh.value()
            ol = self.OLow.value()
            BLAM = self.DetrendMethod.currentText()

            if BLAM == "AccDetrend":
                self.eqs.detrendAcc(oh,ol,curr)
            elif BLAM == "Modified":
                self.eqs.baseLineCorr(oh,ol,curr)
            else:
                pass
            curr = True

        if tadj:
            ntp = self.NTP.value()
            TAM = self.TAMethod.currentText()

            if TAM == "Disp":
                self.eqs.targetDispCorr(ntp,curr)
            elif TAM == "DispVel":
                pass
            elif TAM == "DispVelAcc":
                pass
            else:
                pass
            curr = True

        if not curr:
            self.eqs.a2vd(curr)

        self.plotTH()

    def calSPthenPlotSP(self):
        self.getParaSP()
        self.eqs.getSpectra()
        self.plotSPA()
        self.tabWidget.setCurrentIndex(2)
        self.CDR.setEnabled(True)

    def calASPthenPlotASP(self):
        self.getParaSP()
        self.eqs.getAllSpectra()
        self.plotSP()
        self.tabWidget.setCurrentIndex(3)

    def changeXS_FFT(self):
        X_LOG = self.X_LOG_FFT.isChecked()
        self.ViewFFT.getFigure().add_subplot(1,1,1).set_xscale(XS[X_LOG])
        self.ViewFFT.draw()

    def changeXS_SPA(self):
        X_LOG = self.X_LOG_SPA.isChecked()
        self.ViewSPA.getFigure().add_subplot(1,1,1).set_xscale(XS[X_LOG])
        self.ViewSPA.draw()

    def changeXS_SP(self):
        X_LOG = self.X_LOG_SP.isChecked()
        for ax in self.ViewSP.getFigure().get_axes():
            ax.set_xscale(XS[X_LOG])

        self.ViewSP.draw()

    def plotTH(self):

        raw_on = self.RawON.isChecked()
        fig = self.ViewTH.getFigure()
        fig.clear()
        self.initViewTH()

        accItem, velItem, dispItem = fig.get_axes()

        if raw_on: 
            accItem.plot(self.eqs.t,self.eqs.acc_raw,"0.5",label="Original")
            velItem.plot(self.eqs.t,self.eqs.vel_raw,"0.5",label="Original")
            dispItem.plot(self.eqs.t,self.eqs.disp_raw,"0.5",label="Original")
        
        accItem.plot(self.eqs.t,self.eqs.acc,"b",label="Adjusted")
        velItem.plot(self.eqs.t,self.eqs.vel,"b",label="Adjusted")
        dispItem.plot(self.eqs.t,self.eqs.disp,"b",label="Adjusted")

        self.ViewTH.draw()

        # self.plotFFT()

    def FFTPSD(self):
        i = self.FSP.currentIndex()
        if i == 0:
            self.plotFFT()
        elif i == 1:
            self.plotPSD()
        else:
            pass

    def plotFFT(self):
        self.eqs.fft()
        fig = self.ViewFFT.getFigure()
        fig.clear()
        FFTItem = fig.add_subplot(1,1,1)
        FFTItem.set_title("Fourier Amplitude Spectrum")
        FFTItem.set_xlabel("Frequency (Hz)")
        FFTItem.set_ylabel("Fourier Amplitude")
        n = len(self.eqs.freq)/2
        FFTItem.plot(self.eqs.freq[1:n], 2.0*abs(self.eqs.accf[1:n]), "b")
        self.ViewFFT.draw()

    def plotPSD(self):
        self.eqs.psd()
        fig = self.ViewFFT.getFigure()
        fig.clear()
        FFTItem = fig.add_subplot(1,1,1)
        FFTItem.set_title("Power Spectral Density (Welch's method)")
        FFTItem.set_xlabel("Frequency (Hz)")
        FFTItem.set_ylabel("PSD")
        n = len(self.eqs.freq)/2
        FFTItem.plot(self.eqs.freq[1:n], 2.0*abs(self.eqs.accf[1:n]), "b")
        self.ViewFFT.draw()

    def plotSPA(self):

        fig = self.ViewSPA.getFigure()
        fig.clear()
        self.initViewSPA()
        SPItem = fig.add_subplot(1,1,1)
        if self.DM.currentText() == "linear":
            SPItem.set_xscale(XS[False])
            self.X_LIN_SPA.setChecked(True)
            self.X_LOG_SPA.setChecked(False)
        elif self.DM.currentText() == "log":
            SPItem.set_xscale(XS[True])
            self.X_LIN_SPA.setChecked(False)
            self.X_LOG_SPA.setChecked(True)

        cl = AutoColor()

        for sp in self.eqs.sp:
            if self.SPAPON.isChecked():
                SPItem.plot(sp.SPT,sp.SPA_bak,color="0.5",lw=2)
            if self.SPATON.isChecked():
                SPItem.plot(sp.SPT,sp.SPAT,"--",color=cl.next(),lw=2)
            SPItem.plot(sp.SPT,sp.SPA,color=cl.color,lw=2,label=u"$ \zeta $ = %.2f"%sp.Zeta)
        SPItem.legend(loc=0)
        self.ViewSPA.draw()

    def plotCurrentSPA(self):
        i = self.CDRfit.currentIndex()
        self.plotiSPA(i)

    def plotiSPA(self,i):
        fig = self.ViewSPA.getFigure()
        fig.clear()
        self.initViewSPA()
        SPItem = fig.add_subplot(1,1,1)
        SPItem.set_xscale(XS[self.X_LOG_SPA.isChecked()])

        sp = self.eqs.sp[i]
        if self.SPAPON.isChecked():
            SPItem.plot(sp.SPT,sp.SPA_bak,"0.5",lw=2,label="Pre-Fitted")
        if self.SPATON.isChecked():
            SPItem.plot(sp.SPT,sp.SPAT,"r",lw=2,label="Code Spectrum")
        SPItem.plot(sp.SPT,sp.SPA,"b",lw=2,label="Fitted")
        
        SPItem.legend(loc=0)
        self.ViewSPA.draw()

    def plotSP(self):

        fig = self.ViewSP.getFigure()
        fig.clear()
        self.initViewSP()

        for ax in fig.get_axes():
            xs = self.DM.currentText()
            ax.set_xscale(xs)
            ax.legend(loc=0)
            if xs == "linear":
                self.X_LIN_SP.setChecked(True)
                self.X_LOG_SP.setChecked(False)
            elif xs == "log":
                self.X_LIN_SP.setChecked(False)
                self.X_LOG_SP.setChecked(True)

        SPAItem, SPVItem, SPDItem, SPEItem = fig.get_axes()
        cl = AutoColor()
        for sp in self.eqs.sp:
            SPAItem.plot(sp.SPT,sp.SPA,color=cl.next(),lw=2,label=u"$ \zeta $ = %.2f"%sp.Zeta)
            SPVItem.plot(sp.SPT,sp.SPV,color=cl.color, lw=2,label=u"$ \zeta $ = %.2f"%sp.Zeta)
            SPDItem.plot(sp.SPT,sp.SPD,color=cl.color, lw=2,label=u"$ \zeta $ = %.2f"%sp.Zeta)
            SPEItem.plot(sp.SPT,sp.SPE,color=cl.color, lw=2,label=u"$ \zeta $ = %.2f"%sp.Zeta)
        self.ViewSP.draw()

    def getTXTFileName(self):
        fd = QtGui.QFileDialog()
        self.fileName = fd.getOpenFileName(self, "file", "./data/", "*.txt", "Plain Text File (*.txt *.TXT)")[0]

    def getNGAFileName(self):
        fd = QtGui.QFileDialog()
        self.fileName = fd.getOpenFileName(self, "file", "./data/", "*.at2", "PEER Record File (*.at2 *.AT2)")[0]

    def getParaSP(self):
        Zeta = [float(zi) for zi in self.DR.text().split(",")]
        Tstart = self.TS.value()
        Tstop  = self.TL.value()
        NP = self.NP.value()
        DM = self.DM.currentText()
        SM = self.SM.currentText()
        PD = self.PD.isChecked()
        Tg = self.Tg.value()
        self.eqs.setupSpectra(NP,Tstart,Tstop,DM,Zeta,SM,PD)
        for sp in self.eqs.sp:
            sp.getCodeSpectra(Tg)

    def fitSpectra(self):
        self.getParaSP()
        self.eqs.getSpectra()
        self.eqs.backupSpectra()
        i = self.CDRfit.currentIndex()
        m = self.FitMethod.currentIndex()
        tol = self.Tol.value()
        mit = self.MIT.value()
        
        self.eqs.fitSpectra(i,tol,mit,m)
        self.eqs.getSpectra()
        aer,mer = self.eqs.sp[i].getError()
        self.statusBar().showMessage("Mean Error: %.3f%% \t Max Error: %.3f%%"%(aer*100,mer*100))
        self.plotiSPA(i)
        self.plotTH()

    # def getTHData(self):
    #     Data = spy.vstack((self.eqs.t,self.eqs.acc,self.eqs.vel,self.eqs.disp)).T
    #     return Data

    # def getSPData(self):
    #     i = self.CDR.currentIndex()
    #     Data = spy.vstack((self.eqs.sp[i].SPT,self.eqs.sp[i].SPA,self.eqs.sp[i].SPV,self.eqs.sp[i].SPD,self.eqs.sp[i].SPE)).T
    #     return Data

    # def getSPAData(self):
    #     i = self.CDR.currentIndex()
    #     Data = spy.vstack((self.eqs.sp[i].SPT,self.eqs.sp[i].SPA)).T
    #     return Data