# 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, QCheckBox, QMessageBox, QErrorMessage, \
    QRadioButton

# import TigGUI.kitties.utils

from astropy.io import fits as pyfits

import os.path
import traceback

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

DEG = math.pi / 180

from astLib.astWCS import WCS


class MakeBrickDialog(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("Convert sources to 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="Output FITS file", default_suffix="fits",
                                  file_types="FITS files (*.fits *.FITS)", file_mode=QFileDialog.ExistingFile)
        lo.addWidget(self.wfile)
        # reference frequency
        lo1 = QHBoxLayout()
        lo.addLayout(lo1)
        lo1.setContentsMargins(0, 0, 0, 0)
        label = QLabel("Frequency, MHz:", self)
        lo1.addWidget(label)
        tip = """<P>If your sky model contains spectral information (such as spectral indices), then a brick may be generated
    for a specific frequency. If a frequency is not specified here, the reference frequency of the model sources will be assumed.</P>"""
        self.wfreq = QLineEdit(self)
        self.wfreq.setValidator(QDoubleValidator(self))
        label.setToolTip(tip)
        self.wfreq.setToolTip(tip)
        lo1.addWidget(self.wfreq)
        # beam gain
        lo1 = QHBoxLayout()
        lo.addLayout(lo1)
        lo1.setContentsMargins(0, 0, 0, 0)
        self.wpb_apply = QCheckBox("Apply primary beam expression:", self)
        self.wpb_apply.setChecked(True)
        lo1.addWidget(self.wpb_apply)
        tip = """<P>If this option is specified, a primary power beam gain will be applied to the sources before inserting
    them into the brick. This can be any valid Python expression making use of the variables 'r' (corresponding
    to distance from field centre, in radians) and 'fq' (corresponding to frequency.)</P>"""
        self.wpb_exp = QLineEdit(self)
        self.wpb_apply.setToolTip(tip)
        self.wpb_exp.setToolTip(tip)
        lo1.addWidget(self.wpb_exp)
        # overwrite or add mode
        lo1 = QHBoxLayout()
        lo.addLayout(lo1)
        lo1.setContentsMargins(0, 0, 0, 0)
        self.woverwrite = QRadioButton("overwrite image", self)
        self.woverwrite.setChecked(True)
        lo1.addWidget(self.woverwrite)
        self.waddinto = QRadioButton("add into image", self)
        lo1.addWidget(self.waddinto)
        # add to model
        self.wadd = QCheckBox("Add resulting brick to sky model as a FITS image component", self)
        lo.addWidget(self.wadd)
        lo1 = QHBoxLayout()
        lo.addLayout(lo1)
        lo1.setContentsMargins(0, 0, 0, 0)
        self.wpad = QLineEdit(self)
        self.wpad.setValidator(QDoubleValidator(self))
        self.wpad.setText("1.1")
        lab = QLabel("...with padding factor:", self)
        lab.setToolTip("""<P>The padding factor determines the amount of null padding inserted around the image during
      the prediction stage. Padding alleviates the effects of tapering and detapering in the uv-brick, which can show
      up towards the edges of the image. For a factor of N, the image will be padded out to N times its original size.
      This increases memory use, so if you have no flux at the edges of the image anyway, then a pad factor of 1 is
      perfectly fine.</P>""")
        self.wpad.setToolTip(lab.toolTip())
        self.wadd.toggled[bool].connect(self.wpad.setEnabled)
        self.wadd.toggled[bool].connect(lab.setEnabled)
        self.wpad.setEnabled(False)
        lab.setEnabled(False)
        lo1.addStretch(1)
        lo1.addWidget(lab, 0)
        lo1.addWidget(self.wpad, 1)
        self.wdel = QCheckBox("Remove from the sky model sources that go into the brick", self)
        lo.addWidget(self.wdel)
        # OK/cancel buttons
        lo.addSpacing(10)
        lo2 = QHBoxLayout()
        lo.addLayout(lo2)
        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
        pb = self.model.primaryBeam()
        if pb:
            self.wpb_exp.setText(pb)
        else:
            self.wpb_apply.setChecked(False)
            self.wpb_exp.setText("")
        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
        # read fits file
        busy = BusyIndicator()
        try:
            input_hdu = pyfits.open(filename)[0]
            hdr = input_hdu.header
            # get frequency, if specified
            for axis in range(1, hdr['NAXIS'] + 1):
                if hdr['CTYPE%d' % axis].upper() == 'FREQ':
                    self.wfreq.setText(str(hdr['CRVAL%d' % axis] / 1e+6))
                    break
        except Exception as err:
            busy.reset_cursor()
            self.wfile.setFilename('')
            if not quiet:
                QMessageBox.warning(self, "Error reading FITS",
                                    "Error reading FITS file %s: %s" % (filename, str(err)))
            return None
        self.wokbtn.setEnabled(True)
        # 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) \
                    and os.path.exists(src.shape.filename) and os.path.exists(filename) \
                    and os.path.samefile(src.shape.filename, filename):
                self.wadd.setChecked(True)
                self.wadd.setEnabled(False)
                self.wadd.setText("image already in sky model")
                break
        else:
            self.wadd.setText("add image to sky model")
        busy.reset_cursor()
        return filename

    def accept(self):
        """Tries to make a brick, and closes the dialog if successful."""
        sources = [src for src in self.model.sources if src.selected and src.typecode == 'pnt']
        filename = self.wfile.filename()
        if not self._fileSelected(filename):
            return
        # get PB expression
        pbfunc = None
        if self.wpb_apply.isChecked():
            pbexp = str(self.wpb_exp.text())
            try:
                pbfunc = eval("lambda r,fq:" + pbexp)
            except Exception as err:
                QMessageBox.warning(self, "Error parsing PB experssion",
                                    "Error parsing primary beam expression %s: %s" % (pbexp, str(err)))
                return
        # get frequency
        freq = str(self.wfreq.text())
        freq = float(freq) * 1e+6 if freq else None
        # get pad factor
        pad = str(self.wpad.text())
        pad = max(float(pad), 1) if pad else 1
        # 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
        # reset data if asked to
        if self.woverwrite.isChecked():
            input_hdu.data[...] = 0
        # insert sources
        Imaging.restoreSources(input_hdu, sources, 0, primary_beam=pbfunc, freq=freq)
        # save fits file
        try:
            # pyfits seems to produce an exception:
            #         TypeError: formatwarning() takes exactly 4 arguments (5 given)
            # when attempting to overwrite a file. As a workaround, remove the file first.
            if os.path.exists(filename):
                os.remove(filename)
            input_hdu.writeto(filename)
        except Exception as err:
            traceback.print_exc()
            busy.reset_cursor()
            QMessageBox.warning(self, "Error writing FITS", "Error writing FITS file %s: %s" % (filename, str(err)))
            return
        changed = False
        sources = self.model.sources
        # remove sources from model if asked to
        if self.wdel.isChecked():
            sources = [src for src in sources if not (src.selected and src.typecode == 'pnt')]
            changed = True
        # add image to model if asked to
        if self.wadd.isChecked():
            hdr = input_hdu.header
            # get image parameters
            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
            ra0, dec0 = wcs.pix2wcs(ra0, dec0)
            ra0 *= DEG
            dec0 *= DEG
            sx, sy = wcs.getHalfSizeDeg()
            sx *= DEG
            sy *= DEG
            nx, ny = input_hdu.data.shape[-1:-3:-1]
            # check if this image is already contained in the model
            for src in sources:
                if isinstance(getattr(src, 'shape', None), ModelClasses.FITSImage) and os.path.samefile(
                        src.shape.filename, filename):
                    # update source parameters
                    src.pos.ra, src.pos.dec = ra0, dec0
                    src.flux.I = max_flux
                    src.shape.ex, src.shape.ey = sx, sy
                    src.shape.nx, src.shape.ny = nx, ny
                    src.shape.pad = pad
                    break
            # not contained, make new source object
            else:
                pos = ModelClasses.Position(ra0, dec0)
                flux = ModelClasses.Flux(max_flux)
                shape = ModelClasses.FITSImage(sx, sy, 0, os.path.basename(filename), nx, ny, pad=pad)
                img_src = SkyModel.Source(os.path.splitext(os.path.basename(filename))[0], pos, flux, shape=shape)
                sources.append(img_src)
            changed = True
        if changed:
            self.model.setSources(sources)
            self.model.emitUpdate(SkyModel.SkyModel.UpdateAll, origin=self)
        self.parent().showMessage("Wrote %d sources to FITS file %s" % (len(sources), filename))
        busy.reset_cursor()
        return QDialog.accept(self)


def make_brick(mainwin, model):
    # check that something is selected
    if not [src for src in model.sources if src.selected]:
        mainwin.signalshowErrorMessage.emit(
            "Cannot make FITS brick without a source selection. Please select some sources first.")
        return
    dialog = getattr(mainwin, '_make_brick_dialog', None)
    if not dialog:
        dialog = mainwin._make_brick_dialog = MakeBrickDialog(mainwin)
    dialog.setModel(model)
    # show dialog
    return dialog.exec_()


from TigGUI.Tools import registerTool

registerTool("Make FITS brick from selected sources...", make_brick)
