#-------------------------------------------------------------------------------
# Name:        ApplicationWindow
# Purpose:
#
# Author:      Kelvin
#
# Created:     02/01/2014
#-------------------------------------------------------------------------------

from PyQt4 import QtGui, QtCore
from Skymap import *
from SpectralGraph import *
from ProgressBar import *
from DuchampResults import *

import pyfits
import os

"""Application Window Class Definition"""
# This handles the main window.
class ApplicationWindow(QtGui.QMainWindow):

    # Main Window Initialisation
    def __init__(self):

        # QMainWindow Initialisation
        QtGui.QMainWindow.__init__(self)
        self.setAttribute(QtCore.Qt.WA_DeleteOnClose)
        self.setWindowTitle("Duchamp")

        # File Menu
        self.file_menu = QtGui.QMenu('&File', self)
        self.file_menu.addAction('&Open File', self.openFiles, QtCore.Qt.CTRL + QtCore.Qt.Key_O)
        self.file_menu.addAction('&Save as parameters file', self.saveParameters, QtCore.Qt.CTRL + QtCore.Qt.Key_P)
        self.file_menu.addSeparator()
        self.file_menu.addAction('&Quit', self.close, QtCore.Qt.CTRL + QtCore.Qt.Key_Q)
        self.menuBar().addMenu(self.file_menu)

        # Options Menu
        self.options_menu = QtGui.QMenu('&Options', self)
        self.menuBar().addSeparator()
        self.menuBar().addMenu(self.options_menu)
        self.options_menu.addAction('&Full Screen Mode', self.showFullScreen, QtCore.Qt.CTRL + QtCore.Qt.Key_F)
        self.options_menu.addAction('&Maximized Mode', self.showMaximized, QtCore.Qt.CTRL + QtCore.Qt.Key_M)
        self.options_menu.addAction('&Normal Mode', self.showNormal, QtCore.Qt.CTRL + QtCore.Qt.Key_N)

        # Help Menu
        self.help_menu = QtGui.QMenu('&Help', self)
        self.menuBar().addSeparator()
        self.menuBar().addMenu(self.help_menu)
        self.help_menu.addAction('&About', self.about, QtCore.Qt.CTRL + QtCore.Qt.Key_A)

        # Main Widget
        self.main_widget = QtGui.QWidget(self)
        self.main_widget.setFocus()
        self.setCentralWidget(self.main_widget)
        self.layout = QtGui.QGridLayout(self.main_widget)

        # Main Window Status Bar
        self.statusBar().showMessage("Duchamp Output Window", 2000)

        # Window location and size
        self.center = QtGui.QDesktopWidget().availableGeometry().center()
        self.move(0, 0)
        self.xsize = 2*self.center.x()*0.9
        self.ysize = 2*self.center.y()*0.9
        self.resize(self.xsize, self.ysize)

    # About Message
    def about(self):

        QtGui.QMessageBox.about(self, "About Duchamp",
"""
<i>Duchamp</i> is an object-finder for use
on spectral-line data cubes. The basic execution of <i>Duchamp</i> is to read in a
FITS data cube, find sources in the cube, and produce a text file of positions,
velocities and fluxes of the detections, as well as a postscript file of the spectra
of each detection.

<i>Duchamp</i> has been designed to search for objects in particular sorts of data:
those with relatively small, isolated objects in a large amount of background or
noise. Examples of such data are extragalactic HI surveys, or maser surveys.
<i>Duchamp</i> searches for groups of connected voxels (or pixels) that are all above
some flux threshold. No assumption is made as to the shape of detections, and
the only size constraints applied are those specified by the user.

<i>Duchamp</i> has been written as a three-dimensional source finder, but it is possible
to run it on a two-dimensional image (i.e. one with no frequency or velocity
information), or indeed a one-dimensional array, and many of the features of
the program will work fine. The focus, however, is on object detection in three
dimensions, one of which is a spectral dimension. Note, in particular, that it
does not do any fitting of source profiles, a feature common (and desirable)
for many two-dimensional source finders. This is beyond the current scope of
<i>Duchamp</i>, whose aim is reliable detection of spectral-line objects.

<i>Duchamp</i> provides the ability to pre-process the data prior to searching.
Spectral baselines can be removed, and either smoothing or multi-resolution
wavelet reconstruction can be performed to enhance the completeness and
reliability of the resulting catalogue.
""")

    # This prompts the user to open required files.
    def openFiles(self):

        # Prompt for results file name.
        resultsFileName = QtGui.QFileDialog.getOpenFileName(self, 'Open results File', '.', "text files (*.txt);;All files (*)")

        # If the name is empty, then we did not load any files.
        if not resultsFileName:
            QtGui.QMessageBox.information(self, 'Message', 'No file loaded.', QtGui.QMessageBox.Ok)
            return

        # Open up the results file.
        resultsFile = open(resultsFileName)

        # Get the first line of the results file.
        for line in resultsFile:
            firstline = line
            break

        # Determine if the content of the results file is correct.
        if "# Results of the Duchamp source finder" not in firstline:
            QtGui.QMessageBox.information(self, 'Message', 'Incorrect content for results file. No file loaded.', QtGui.QMessageBox.Ok)
            return

        # Find the maskfits and fits file name.
        maskfitsFileName = ''
        fitsFileName = ''
        for line in resultsFile:

            if 'imageFile' in line and ' = ' in line:

                fitsFileName = line.split(' = ')[-1].split()[0].strip()

            if 'flagOutputMask' in line and '>' in line:

                maskfitsFileName = line.split('>')[-1].split()[0].strip()

        # If we cannot find the fits file name, prompt the user to find it.
        if not fitsFileName:
            QtGui.QMessageBox.information(self, 'Message', 'Results file do not contain input "imageFile" name. Please specify the original input fits file used in source finding.', QtGui.QMessageBox.Ok)
            
            # Prompt for the input cube fits file name.
            fitsFileName = QtGui.QFileDialog.getOpenFileName(self, 'Open input File', '.', "fits files (*.fits);;All files (*)")

            # If the name is empty, then we did not load any files.
            if not fitsFileName:
                QtGui.QMessageBox.information(self, 'Message', 'No file loaded.', QtGui.QMessageBox.Ok)
                return

        # If the file found does not exist, then let the user know and give the user a final chance.
        while not os.path.isfile(fitsFileName):
            QtGui.QMessageBox.information(self, 'Message', 'Input "imageFile" <%s> specified in results file cannot be found in current directory. Please specify the original input fits file used in source finding.'%fitsFileName, QtGui.QMessageBox.Ok)
            
            # Prompt for the input cube fits file name.
            fitsFileName = QtGui.QFileDialog.getOpenFileName(self, 'Open input File', '.', "fits files (*.fits);;All files (*)")

            # If the user cancelled the loading, then do not prompt the user anymore.
            if not fitsFileName:
                return

        # If we cannot find the maskfits file name, prompt the user to find it.
        if not maskfitsFileName:
            QtGui.QMessageBox.information(self, 'Message', 'Results file do not contain MASK fits file name. Please specify MASK fits file.', QtGui.QMessageBox.Ok)

            # Prompt for the mask cube fits file name.
            maskfitsFileName = QtGui.QFileDialog.getOpenFileName(self, 'Open mask File', '.', "MASK fits files (*.MASK.fits);;fits files (*.fits);;All files (*)")
        
            # If the name is empty, then we did not load any files.
            if not maskfitsFileName:
                QtGui.QMessageBox.information(self, 'Message', 'No file loaded.', QtGui.QMessageBox.Ok)
                return

        # If the file found does not exist, then let the user know and give the user a final chance.
        while not os.path.isfile(maskfitsFileName):
            QtGui.QMessageBox.information(self, 'Message', 'MASK fits file <%s> specified in results file cannot be found in current directory. Please specify the MASK fits file.'%maskfitsFileName, QtGui.QMessageBox.Ok)
            
            # Prompt for the mask cube fits file name.
            maskfitsFileName = QtGui.QFileDialog.getOpenFileName(self, 'Open mask File', '.', "MASK fits files (*.MASK.fits);;fits files (*.fits);;All files (*)")
        
            # If the user cancelled the loading, then do not prompt the user anymore.
            if not maskfitsFileName:
                return

        # Store the results from the results file as a class
        self.Results = DuchampResults(resultsFileName)

        # Set the names of the files
        self.Results.maskfitsFileName = maskfitsFileName
        self.Results.fitsFileName = fitsFileName

        # If the file names are okay, then check and load the files using those names.
        self.checkAndLoadFiles(maskfitsFileName, fitsFileName)

    # This makes sure that the files are correct and in the correct format.
    # If everything is good, it will start the software.
    def checkAndLoadFiles(self, maskfitsFileName, fitsFileName):
        
        # Try reading the files.
        try:
            # Open the mask cube fits file.
            maskFits = pyfits.open(maskfitsFileName)
        
            # If the 'BUNIT' option in the mask cube fits header is not 'Object ID', then the program will not start.
            if maskFits[0].header.get('BUNIT') != 'Object ID':
                QtGui.QMessageBox.warning(self, 'Incorrect specification in  .MASK.fits file', 'The "BUNIT" option in "%s" is not set to "Object ID". Files not loaded.'%maskfitsFileName, QtGui.QMessageBox.Ok)
                return(0)

            # Open the input cube fits file.
            realFits = pyfits.open(fitsFileName)

            # If the two cubes do not have matching shapes, then they are obviously not meant for each other. The program will not start.
            if maskFits[0].data.shape != realFits[0].data.shape:
                QtGui.QMessageBox.warning(self, 'Cube sizes mismatch', 'Fits cube sizes differ in shape in files "%s" and "%s". Files not loaded.'%(maskfitsFileName, fitsFileName), QtGui.QMessageBox.Ok)
                return(0)

            # If everything is fine, we will confirm the choice of files with the user.
            SuccessMsg = QtGui.QMessageBox()
            SuccessMsg.setWindowTitle("Files Loaded")
            SuccessMsg.setText('Files\n"%s" and \n"%s"\n are loaded successfully!'%(maskfitsFileName, fitsFileName))
            SuccessMsg.setInformativeText('Please confirm that those are the correct files.')
            SuccessMsg.setStandardButtons(QtGui.QMessageBox.Yes | QtGui.QMessageBox.Discard)
            SuccessMsg.setDefaultButton(QtGui.QMessageBox.Yes)

            # If the user is fine with the choice, then start the program with those files.
            if SuccessMsg.exec_() == QtGui.QMessageBox.Yes:

                # This is the name of the input file
                self.fitsName = fitsFileName.split('/')[-1]

                # This is the name of the maskfits file
                self.maskfitsName = maskfitsFileName.split('/')[-1]

                # Those are the fits data
                self.maskFits = maskFits
                self.realFits = realFits

                # Now we can start processing the data
                self.processData()
                return(1)

            # Otherwise, discard the files.
            QtGui.QMessageBox.information(self, 'Message', 'Files discarded.', QtGui.QMessageBox.Ok)
            return(0)
     
        # If anything went wrong in the above process, then file reading has failed.
        except Exception:
            QtGui.QMessageBox.warning(self, 'Error', 'Reading of fits files failed. Files not loaded.', QtGui.QMessageBox.Ok)

        return(0)

    # This loads the required data from the files that contains the data.
    def processData(self):

        # Use a progress bar to indicate the loading process
        self.pbar = ProgressBar(title = "Processing data", total = 100)
        self.pbar.show()

        # Obtain the two cubes
        self.maskCube = self.maskFits[0].data
        self.realCube = self.realFits[0].data

        # Those are the size of our data cubes
        [self.lenZ, self.lenY, self.lenX] = self.maskCube.shape

        # This is the cube used to produce the final image
        self.finalCube = zeros([self.lenZ, self.lenY, self.lenX])

        # Those are our mask image, real image, and final image respectively
        self.maskImage = zeros([self.lenY, self.lenX])
        self.realImage = zeros([self.lenY, self.lenX])
        self.finalImage = zeros([self.lenY, self.lenX])

        # This will contain all the moment maps with their index matching its object ID
        self.momentMaps = ['Moment Maps']

        # Here we compute the final cube from the mask cube and real cube

        # Go through each object we found
        for ObjID in range(1, self.Results.totalDetections + 1):

            # Obtain information regarding that source
            ObjStats = self.Results.SkyObjects[ObjID].stats

            # Obtain the spatial and spectral bound of that source
            x1 = ObjStats['X1']
            x2 = ObjStats['X2']
            y1 = ObjStats['Y1']
            y2 = ObjStats['Y2']
            z1 = ObjStats['Z1']
            z2 = ObjStats['Z2']

            # Calculate the moment map of the source as well as add the source to the final cube
            momentMap = zeros([self.lenY, self.lenX])

            for z in range(z1, z2 + 1):

                momentMap += self.realCube[z]

                for y in range(y1, y2 + 1):
                    for x in range(x1, x2 + 1):

                        if self.maskCube[z][y][x] == ObjID:
                            self.finalCube[z][y][x] = self.realCube[z][y][x]

            # Add this moment map to the list of moment maps
            self.momentMaps += [momentMap]

            # And also compute the mask image from the mask cube
            for z in range(self.lenZ):
                self.maskImage[self.maskCube[z] == ObjID] = ObjID

            # Update the progress bar accordingly
            self.pbar.update_progressbar(100*ObjID/self.Results.totalDetections + 1)

        # Compute the real image from the real cube
        for z in range(self.lenZ):
            self.realImage += self.realCube[z]

        # Compute the final image from the final cube
        for z in range(self.lenZ):
            self.finalImage += self.finalCube[z]
        
        # The data processing has finished
        self.pbar.close()

        # The program can now start
        self.start()

    # This defines what happens when the software starts.
    def start(self):

        # Use a progress bar to indicate the start up progress
        self.pbar = ProgressBar(title = "Loading Skymap", total = 100)
        self.pbar.show()

        # Start the skymap.
        self.startSkymap()
        self.pbar.update_progressbar(26)

        # Start the spectral graph.
        self.startSpectralGraph()
        self.pbar.update_progressbar(51)

        # Start the zoom graph
        self.startZoomGraph()
        self.pbar.update_progressbar(76)

        # Start the moment map
        self.startMomentMap()
        self.pbar.update_progressbar(100)

        self.pbar.close()

    # This starts the skymap.
    def startSkymap(self):

        # Initialise the skymap.
        self.skymap = Skymap(self.Results, maskImage = self.maskImage, finalImage = self.finalImage)
        self.skymap.drawImage()
        self.skymap.enableInteractivity()
        self.skymap.fig.suptitle(self.fitsName)

        # Determine its size and location.
        self.skymap.setMinimumHeight(0.5*self.ysize)
        self.skymap.setMinimumWidth(0.3*self.xsize)
        self.layout.addWidget(self.skymap, 0, 0, 4, 4)

        # Connect the signal from skymap to update the spectral graph.
        self.skymap.graphSignal.connect(self.updateSpectralGraph)
        self.skymap.graphSignal.connect(self.updateZoomGraph)
        self.skymap.graphSignal.connect(self.updateMomentMap)

    # This starts the spectral graph.
    def startSpectralGraph(self):

        # Initialise the spectral graph.
        self.spectralGraph = SpectralGraph(self.Results, self.maskFits, self.realFits, self.skymap.maskImage)

        # Determine its size and location.
        self.spectralGraph.setMinimumHeight(0.3*self.ysize)
        self.spectralGraph.setMinimumWidth(self.xsize)

        self.layout.addWidget(self.spectralGraph, 4, 0, 2, 8)

        # Initialise the radio buttons for display options.
        self.displayMaxButton = QtGui.QRadioButton("Peak Detected Intensity")
        self.displaySumButton = QtGui.QRadioButton("Total Detected Intensity")

        # Determine its location.
        self.layout.addWidget(self.displayMaxButton, 0, 4, 1, 1)
        self.layout.addWidget(self.displaySumButton, 0, 5, 1, 1)
        
        # Determine its functionality.
        self.displayMaxButton.toggled.connect(self.spectralGraph.plotMax)
        self.displaySumButton.toggled.connect(self.spectralGraph.plotSum)
        self.displayMaxButton.setChecked(True)


    # This updates the spectral graph.
    def updateSpectralGraph(self):

        # Set the current object to plot the spectral graph with as the highlighted object in the skymap, and plot the spectral graph.
        self.spectralGraph.setCurrentObjID(self.skymap.highlightedObj)
        self.spectralGraph.plotSource()

    # This starts the zoom graph.
    def startZoomGraph(self):

        # Initialise the zoom graph
        self.zoomGraph = ZoomGraph(self.Results, self.maskFits, self.realFits, self.skymap.maskImage)

        # Determine its location
        self.layout.addWidget(self.zoomGraph, 1, 4, 3, 2)
        
        # Link the radio button functionalities to the zoom graph
        self.displayMaxButton.toggled.connect(self.zoomGraph.plotMax)
        self.displaySumButton.toggled.connect(self.zoomGraph.plotSum)

    # This updates the zoom graph.
    def updateZoomGraph(self):

        # Set the current object to plot the zoom graph with as the highlighted object in the skymap, and plot the zoom graph.
        self.zoomGraph.setCurrentObjID(self.skymap.highlightedObj)
        self.zoomGraph.plotSource()

    # This starts the moment map.
    def startMomentMap(self):

        # Initialise the moment map
        self.momentMap = MomentMap(self.Results, maskImage = self.maskImage, finalImage = self.finalImage, momentMaps = self.momentMaps)

        # Determine its location
        self.layout.addWidget(self.momentMap, 1, 6, 3, 2)

    # This updates the moment map.
    def updateMomentMap(self):

        # Set the current object to draw the moment map with as the highlighted object in the skymap, and draw the moment map
        self.momentMap.setCurrentObjID(self.skymap.highlightedObj)
        self.momentMap.drawImage()

    # This saves the parameters used to generate this output into a parameter file.
    def saveParameters(self):

        # Prompt for the name of the file.
        parameterFilePath = QtGui.QFileDialog.getSaveFileName(self, 'Save parameters as', '.', "text files (*.txt);;All files (*)")

        # Write the parameter.
        try:
            writeParamFile(parameterFilePath, self.Results.parameters, self.Results.parameterName)

        except Exceptions:
            pass