# Copyright (C) 2002-2022
# The MeqTree Foundation &
# ASTRON (Netherlands Foundation for Research in Astronomy)
# P.O.Box 2, 7990 AA Dwingeloo, The Netherlands
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>,
# or write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#

import math
from PyQt5.QtWidgets import *

from PyQt5.Qt import QObject, QHBoxLayout, QFileDialog, pyqtSignal, QLabel, \
    QLineEdit, QDialog, QDoubleValidator, QVBoxLayout, \
    QPushButton, Qt, QGridLayout, QMessageBox, QErrorMessage

import TigGUI.kitties.utils

from astropy.io import fits as pyfits

import os.path

from TigGUI.kitties.widgets import BusyIndicator
from TigGUI.Widgets import FileSelector
from Tigger.Models import SkyModel, ModelClasses

DEG = math.pi / 180

from astLib.astWCS import WCS


class AddBrickDialog(QDialog):
    def __init__(self, parent, modal=True, flags=Qt.WindowFlags()):
        QDialog.__init__(self, parent, flags)
        self.model = None
        self._model_dir = None
        self.setModal(modal)
        self.setWindowTitle("Add FITS brick")
        lo = QVBoxLayout(self)
        lo.setContentsMargins(10, 10, 10, 10)
        lo.setSpacing(5)
        # file selector
        self.wfile = FileSelector(self, label="FITS filename:", dialog_label="FITS file", default_suffix="fits",
                                  file_types="FITS files (*.fits *.FITS)", file_mode=QFileDialog.ExistingFile)
        lo.addWidget(self.wfile)
        # overwrite or add mode
        lo1 = QGridLayout()
        lo.addLayout(lo1)
        lo1.setContentsMargins(0, 0, 0, 0)
        lo1.addWidget(QLabel("Padding factor:", self), 0, 0)
        self.wpad = QLineEdit("2", self)
        self.wpad.setValidator(QDoubleValidator(self))
        lo1.addWidget(self.wpad, 0, 1)
        lo1.addWidget(QLabel("Assign source name:", self), 1, 0)
        self.wname = QLineEdit(self)
        lo1.addWidget(self.wname, 1, 1)
        # OK/cancel buttons
        lo.addSpacing(10)
        lo2 = QHBoxLayout()
        lo.addLayout(lo2)
        lo2.setContentsMargins(0, 0, 0, 0)
        lo2.setContentsMargins(5, 5, 5, 5)
        self.wokbtn = QPushButton("OK", self)
        self.wokbtn.setMinimumWidth(128)
        self.wokbtn.clicked.connect(self.accept)
        self.wokbtn.setEnabled(False)
        cancelbtn = QPushButton("Cancel", self)
        cancelbtn.setMinimumWidth(128)
        cancelbtn.clicked.connect(self.reject)
        lo2.addWidget(self.wokbtn)
        lo2.addStretch(1)
        lo2.addWidget(cancelbtn)
        self.setMinimumWidth(384)
        # signals
        self.wfile.filenameSelected.connect(self._fileSelected)
        # internal state
        self.qerrmsg = QErrorMessage(self)

    def setModel(self, model):
        self.model = model
        if model.filename():
            self._model_dir = os.path.dirname(os.path.abspath(model.filename()))
        else:
            self._model_dir = os.path.abspath('.')
        self.wfile.setDirectory(self._model_dir)
        self._fileSelected(self.wfile.filename(), quiet=True)

    def _fileSelected(self, filename, quiet=False):
        self.wokbtn.setEnabled(False)
        if not filename:
            return None
        # check that filename matches model
        if not os.path.samefile(self._model_dir, os.path.dirname(filename)):
            self.wfile.setFilename('')
            if not quiet:
                QMessageBox.warning(self, "Directory mismatch", """<P>The FITS file must reside in the same directory
          as the current sky model.</P>""")
            self.wfile.setDirectory(self._model_dir)
            return None
        # if filename is not in model already, enable the "add to model" control
        for src in self.model.sources:
            if isinstance(getattr(src, 'shape', None), ModelClasses.FITSImage):
                if os.path.exists(src.shape.filename) and os.path.samefile(src.shape.filename, filename):
                    if not quiet:
                        QMessageBox.warning(self, "Already in model",
                                            "This FITS brick is already present in the model.")
                    self.wfile.setFilename('')
                    return None
        if not str(self.wname.text()):
            self.wname.setText(os.path.splitext(os.path.basename(str(filename)))[0])
        self.wokbtn.setEnabled(True)
        return filename

    def accept(self):
        """Tries to add brick, and closes the dialog if successful."""
        filename = self.wfile.filename()
        # read fits file
        busy = BusyIndicator()
        try:
            input_hdu = pyfits.open(filename)[0]
        except Exception as err:
            busy.reset_cursor()
            QMessageBox.warning(self, "Error reading FITS", "Error reading FITS file %s: %s" % (filename, str(err)))
            return
        # check name
        srcname = str(self.wname.text()) or os.path.splitext(os.path.basename(str(filename)))[0]
        if srcname in set([src.name for src in self.model.sources]):
            QMessageBox.warning(self, "Already in model",
                                "<p>The model already contains a source named '%s'. Please select a different name.</p>" % srcname)
            return
        # get image parameters
        hdr = input_hdu.header
        max_flux = float(input_hdu.data.max())
        wcs = WCS(hdr, mode='pyfits')
        # Get reference pixel coordinates
        # wcs.getCentreWCSCoords() doesn't work, as that gives us the middle of the image
        # So scan the header to get the CRPIX values
        ra0 = dec0 = 1
        for iaxis in range(hdr['NAXIS']):
            axs = str(iaxis + 1)
            name = hdr.get('CTYPE' + axs, axs).upper()
            if name.startswith("RA"):
                ra0 = hdr.get('CRPIX' + axs, 1) - 1
            elif name.startswith("DEC"):
                dec0 = hdr.get('CRPIX' + axs, 1) - 1
        # convert pixel to degrees
        #    print ra0,dec0
        ra0, dec0 = wcs.pix2wcs(ra0, dec0)
        ra0 *= DEG
        dec0 *= DEG
        #    print ModelClasses.Position.ra_hms_static(ra0)
        #    print ModelClasses.Position.dec_sdms_static(dec0)
        sx, sy = wcs.getHalfSizeDeg()
        sx *= DEG
        sy *= DEG
        nx, ny = input_hdu.data.shape[-1:-3:-1]
        pos = ModelClasses.Position(ra0, dec0)
        flux = ModelClasses.Flux(max_flux)
        shape = ModelClasses.FITSImage(sx, sy, 0, os.path.basename(filename), nx, ny,
                                       pad=float(str(self.wpad.text()) or "1"))
        img_src = SkyModel.Source(srcname, pos, flux, shape=shape)
        self.model.setSources(self.model.sources + [img_src])
        self.model.emitUpdate(SkyModel.SkyModel.UpdateAll, origin=self)
        busy.reset_cursor()
        return QDialog.accept(self)


def add_brick(mainwin, model):
    dialog = getattr(mainwin, '_add_brick_dialog', None)
    if not dialog:
        dialog = mainwin._add_brick_dialog = AddBrickDialog(mainwin)
    dialog.setModel(model)
    # show dialog
    return dialog.exec_()


from TigGUI.Tools import registerTool

registerTool("Add FITS brick to model...", add_brick)
