# 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

import os.path
import time
from PyQt5.Qt import QObject
from PyQt5.QtCore import pyqtSignal
from scipy.ndimage import measurements
import numpy as np

import TigGUI.kitties.utils
from TigGUI.Images.Colormaps import HistEqIntensityMap, LogIntensityMap, CubeHelixColormap
from TigGUI.kitties.widgets import BusyIndicator

_verbosity = TigGUI.kitties.utils.verbosity(name="rc")
dprint = _verbosity.dprint
dprintf = _verbosity.dprintf

from TigGUI.Images import Colormaps

import TigGUI.kitties.config


ImageConfigFile = TigGUI.kitties.config.DualConfigParser("tigger.images.conf")


class RenderControl(QObject):
    """RenderControl represents all the options (slices, color and intensity policy data) associated with an image. This object is shared by various GUI elements
    that control the rendering of images.
    """
    intensityMapChanged = pyqtSignal(object, float)
    colorMapChanged = pyqtSignal(object)
    dataSubsetChanged = pyqtSignal(np.ndarray, tuple, str, str)
    displayRangeChanged = pyqtSignal([float, float], [np.float32, np.float32], [HistEqIntensityMap, float])  # on file save np.float32's become float's on reload?
    displayRangeLocked = pyqtSignal(bool)

    SUBSET_FULL = "full"
    SUBSET_SLICE = "slice"
    SUBSET_RECT = "rect"

    def __init__(self, image, parent):
        QObject.__init__(self, parent)
        self.image = image
        self._config = TigGUI.kitties.config.SectionParser(ImageConfigFile, os.path.normpath(
            os.path.abspath(image.filename))) if image.filename else None
        # figure out the slicing -- find extra axes with size > 1
        # self._current_slice contains all extra axis, including the size-1 ones
        # self._sliced_axes is a list of (iextra,axisname,labels) tuples for size>1 axes
        # where iextra is an index into self._current_slice.
        self._current_slice = [0] * image.numExtraAxes()
        self._slice_dims = [1] * image.numExtraAxes()
        self._sliced_axes = []
        for i in range(image.numExtraAxes()):
            iaxis, axisname, labels = image.extraAxisNumberNameLabels(i)
            self._slice_dims[i] = len(labels)
            if len(labels) > 1:
                self._sliced_axes.append((i, axisname, labels))
        # set the full image range (i.e. mix/max) and current slice range
        dprint(2, "getting data min/max")
        self._fullrange = self._slicerange = image.dataMinMax()[:2]
        dprint(2, "done")
        # create dict of intensity maps
        log_cycles = self._config.getfloat("intensity-log-cycles", 6) if self._config else 6
        self._imap_list = (
            ('Linear', Colormaps.LinearIntensityMap()),
            ('Histogram-equalized', Colormaps.HistEqIntensityMap()),
            ('log(val-min)', Colormaps.LogIntensityMap(log_cycles))
        )
        # create list of color maps
        self._cmap_list = Colormaps.getColormapList()
        default_cmap = 0
        for i, cmap in enumerate(self._cmap_list):
            if isinstance(cmap, Colormaps.ColormapWithControls):
                if self._config:
                    cmap.loadConfig(self._config)
                cmap.colormapChanged.connect(self.updateColorMapParameters)
            if isinstance(cmap, Colormaps.CubeHelixColormap):
                default_cmap = i
        # set the initial intensity map
        imap = self._config.getint("intensity-map-number", 0) if self._config else 0
        cmap = self._config.getint("colour-map-number", default_cmap) if self._config else default_cmap
        imap = max(min(len(self._imap_list) - 1, imap), 0)
        cmap = max(min(len(self._cmap_list) - 1, cmap), 0)
        self._current_imap_index = imap
        self._current_cmap_index = cmap
        self.image.setIntensityMap(self._imap_list[imap][1])
        self.image.setColorMap(self._cmap_list[cmap])

        # cache of min/max values for each slice, as these can be slowish to recompute when flipping slices
        self._sliceranges = {}
        # This is the data subset corresponding to the current display range. When the display range is set to
        # _fullrange, this is the image cube. When it is set to _slicerange, this is the current image slice. When
        # setLMRectDisplayRange() or setWindowDisplayRange() is used to set the range to the specified window,
        # this is the a subset of the current slice. The data subset is passed to setDataSubset() of the intensity mapper object
        self._displaydata = None
        # This is a tuple of the extrema of the current data subset. This is not quite the same thing as self._displayrange below.
        # When the display range is reset to cube/slice/window, _displayrange is set to _displaydata_minmax. But if
        # setDisplayRange() is subsequently called (e.g. if the user manually enters new values into the Range boxes), then
        # _displayrange will be set to something else until the next reset....() call.
        self._displaydata_minmax = None
        # This is a low,high tuple of the current display range -- will be initialized by resetFullDisplayRange()
        self._displayrange = None
        if self._config and self._config.has_option("range-min") and self._config.has_option("range-max"):
            display_range = self._config.getfloat("range-min"), self._config.getfloat("range-max")
        else:
            display_range = None
        self.setFullSubset(display_range, write_config=False)
        # setup initial slice
        if self.hasSlicing():
            if self._config and self._config.has_option("slice"):
                try:
                    curslice = list(map(int, self._config.get("slice").split()))
                except:
                    curslice = []
                if len(curslice) == len(self._current_slice):
                    for iaxis, i in enumerate(curslice):
                        naxis = len(self.image.extraAxisValues(iaxis))
                        i = min(naxis - 1, max(0, i))
                        self._current_slice[iaxis] = i
            self.selectSlice(self._current_slice, write_config=False)
        # lock display range if so configured
        self._lock_display_range = self._config.getbool("lock-range", 0) if self._config else False
        if self._lock_display_range:
            self.lockDisplayRange(True, write_config=False)

    def startSavingConfig(self, image_filename):
        """Saves the current configuration under the specified image filename"""
        self._config = TigGUI.kitties.config.SectionParser(ImageConfigFile, os.path.normpath(os.path.abspath(image_filename)))
        if self._displayrange:
            self._config.set("range-min", self._displayrange[0], save=False)
            self._config.set("range-max", self._displayrange[1], save=False)
        if self._current_slice:
            self._config.set("slice", " ".join(map(str, self._current_slice)), save=False)
        for cmap in self._cmap_list:
            if isinstance(cmap, Colormaps.ColormapWithControls):
                cmap.saveConfig(self._config, save=False)
        self._config.set("intensity-map-number", self._current_imap_index, save=False)
        self._config.set("colour-map-number", self._current_cmap_index, save=False)
        self._config.set("lock-range", self._lock_display_range, save=True)

    def hasSlicing(self):
        """Returns True if image is a cube, and so has non-trivial slicing axes"""
        return bool(self._sliced_axes)

    def slicedAxes(self):
        """Returns list of (axis_num,name,label_list) tuples per each non-trivial slicing axis"""
        return self._sliced_axes

    def incrementSlice(self, iaxis, incr, write_config=True):
        dprint(2, "incrementing slice axis", iaxis, "by", incr)
        self._current_slice[iaxis] = (self._current_slice[iaxis] + incr) % self._slice_dims[iaxis]
        self._updateSlice(write_config)

    def changeSlice(self, iaxis, index, write_config=True):
        dprint(2, "changing slice axis", iaxis, "to", index)
        if self._current_slice[iaxis] != index:
            self._current_slice[iaxis] = index
            self._updateSlice(write_config)

    def selectSlice(self, indices, write_config=True):
        """Selects slice given by indices"""
        dprint(2, "selecting slice", indices)
        self._current_slice = list(indices)
        self._updateSlice(write_config)

    def _updateSlice(self, write_config=True):
        """Common internal method called to finalize changes to _current_slice"""
        busy = BusyIndicator()
        dprint(2, "_updateSlice", self._current_slice, time.time() % 60)
        indices = tuple(self._current_slice)
        self.image.selectSlice(*indices)
        dprint(2, "image slice selected", time.time() % 60)
        img = self.image.image()
        self._slicerange = self._sliceranges.get(indices)
        if self._slicerange is None:
            self._slicerange = self._sliceranges[indices] = self.image.imageMinMax()[:2]
        dprint(2, "min/max updated", time.time() % 60)
        self.setSliceSubset(set_display_range=False)
        if write_config and self._config:
            self._config.set("slice", " ".join(map(str, indices)))
        busy.reset_cursor()

    def displayRange(self):
        return self._displayrange

    def currentSlice(self):
        return self._current_slice

    def sliceDimensions(self):
        return self._slice_dims

    def getIntensityMapNames(self):
        return [name for name, imap in self._imap_list]

    def currentIntensityMapNumber(self):
        return self._current_imap_index

    def currentIntensityMap(self):
        return self.image.intensityMap()

    def setIntensityMapNumber(self, index, write_config=True):
        busy = BusyIndicator()
        self._current_imap_index = index
        imap = self._imap_list[index][1]
        imap.setDataSubset(self._displaydata, self._displaydata_minmax)
        imap.setDataRange(*self._displayrange)
        self.image.setIntensityMap(imap)
        self.intensityMapChanged.emit(imap, index)
        if self._config and write_config:
            self._config.set("intensity-map-number", index)
        busy.reset_cursor()

    def setIntensityMapLogCycles(self, cycles, notify_image=True, write_config=True):
        busy = BusyIndicator()
        imap = self.currentIntensityMap()
        if isinstance(imap, Colormaps.LogIntensityMap):
            imap.log_cycles = cycles
            if notify_image:
                self.image.setIntensityMap()
            self.intensityMapChanged.emit(imap, self._current_imap_index)
        if self._config and write_config:
            self._config.set("intensity-log-cycles", cycles)
        busy.reset_cursor()

    def lockDisplayRangeForAxis(self, iaxis, lock):
        pass

    def getColormapList(self):
        return self._cmap_list

    def updateColorMapParameters(self):
        """Call this when the colormap parameters have changed"""
        busy = BusyIndicator()
        self.image.updateCurrentColorMap()
        if self._config:
            self._cmap_list[self._current_cmap_index].saveConfig(self._config)
        busy.reset_cursor()

    def setColorMapNumber(self, index, write_config=True):
        busy = BusyIndicator()
        self._current_cmap_index = index
        cmap = self._cmap_list[index]
        self.image.setColorMap(cmap)
        self.colorMapChanged.emit(cmap)
        if self._config and write_config:
            self._config.set("colour-map-number", index)
        busy.reset_cursor()

    def currentSubset(self):
        """Returns tuple of subset,(dmin,dmax),description for current data subset"""
        return self._displaydata, self._displaydata_minmax, self._displaydata_desc, self._displaydata_type

    def _resetDisplaySubset(self, subset, desc, range=None, set_display_range=True, write_config=True,
                            subset_type=None):
        dprint(4, "setting display subset")
        self._displaydata = subset
        self._displaydata_desc = desc
        self._displaydata_minmax = range = range or measurements.extrema(subset)[:2]
        self._displaydata_type = subset_type
        dprint(4, "range set")
        self.image.intensityMap().setDataSubset(self._displaydata, minmax=range)
        self.image.setIntensityMap(emit=False)
        dprint(2, f"dataSubsetChanged {type(subset)}, {type(range)}, {type(desc)}, {type(subset_type)}")
        self.dataSubsetChanged.emit(subset, range, desc, subset_type)
        if set_display_range:
            self.setDisplayRange(write_config=write_config, *range)

    def setFullSubset(self, display_range=None, write_config=True):
        shapedesc = "\u00D7".join(["%d" % x for x in
                                    list(self.image.imageDims()) + [len(labels) for iaxis, name, labels in
                                                                    self._sliced_axes]])
        desc = "full cube" if self._sliced_axes else "full image"
        self._resetDisplaySubset(self.image.data(), desc, range=self._fullrange, subset_type=self.SUBSET_FULL,
                                 write_config=write_config, set_display_range=False)
        self.setDisplayRange(write_config=write_config, *(display_range or self._fullrange))

    def _makeSliceDesc(self):
        """Makes a description of the current slice"""
        if not self._sliced_axes:
            return "full image"
        descs = []
        for iextra, name, labels in self._sliced_axes:
            if name.upper() not in ["STOKES", "COMPLEX"]:
                descs.append("%s=%s" % (name, labels[self._current_slice[iextra]]))
            else:
                descs.append(labels[self._current_slice[iextra]])
        return "%s plane" % (" ".join(descs),)

    def setSliceSubset(self, set_display_range=True, write_config=True): \
            return self._resetDisplaySubset(self.image.image(), self._makeSliceDesc(), self._slicerange,
                                            subset_type=self.SUBSET_SLICE,
                                            set_display_range=set_display_range, write_config=write_config)

    def _setRectangularSubset(self, xx1, xx2, yy1, yy2):
        descs = []
        nx, ny = self.image.imageDims()
        if xx1 or xx2 != nx:
            descs.append("x=%d:%d" % (xx1, xx2))
        if yy1 or yy2 != ny:
            descs.append("y=%d:%d" % (yy1, yy2))
        if descs:
            descs.append("in")
        descs.append(self._makeSliceDesc())
        return self._resetDisplaySubset(self.image.image()[xx1:xx2, yy1:yy2], " ".join(descs),
                                        subset_type=self.SUBSET_RECT)

    def _lmRectToPix(self, rect):
        """helper function -- converts an LM rectangle to pixel coordinates"""
        if rect.width() and rect.height():
            # convert to pixel coordinates
            x1, y1, x2, y2 = rect.getCoords()
            x1, y1 = self.image.lmToPix(x1, y1)
            x2, y2 = self.image.lmToPix(x2, y2)
            dprint(2, x1, y1, x2, y2)
            xx1, xx2 = int(math.floor(min(x1, x2))), int(math.ceil(max(x1, x2)))
            yy1, yy2 = int(math.floor(min(y1, y2))), int(math.ceil(max(y1, y2)))
            dprint(2, xx1, yy1, xx2, yy2)
            # ensure limits
            nx, ny = self.image.imageDims()
            xx1, xx2 = max(xx1, 0), min(xx2, nx)
            yy1, yy2 = max(yy1, 0), min(yy2, ny)
            dprint(2, xx1, yy1, xx2, yy2)
            # check that we actually selected some valid pixels
            if xx1 < xx2 and yy1 < yy2:
                return xx1, xx2, yy1, yy2
        return None, None, None, None

    def setLMRectSubset(self, rect):
        xx1, xx2, yy1, yy2 = self._lmRectToPix(rect)
        if xx1 is not None:
            return self._setRectangularSubset(xx1, xx2, yy1, yy2)

    def getLMRectStats(self, rect):
        xx1, xx2, yy1, yy2 = self._lmRectToPix(rect)
        if xx1 is not None:
            subset = self.image.image()[xx1:xx2, yy1:yy2]
            subset, mask = self.image.optimalRavel(subset)
            mmin, mmax = measurements.extrema(subset, labels=mask, index=None if mask is None else False)[:2]
            mean = measurements.mean(subset, labels=mask, index=None if mask is None else False)
            std = measurements.standard_deviation(subset, labels=mask, index=None if mask is None else False)
            ssum = measurements.sum(subset, labels=mask, index=None if mask is None else False)
            return xx1, xx2, yy1, yy2, mmin, mmax, mean, std, ssum, subset.size
        return None

    def setWindowSubset(self, rect=None):
        rect = rect or self.image.currentRectPix()
        if rect.width() and rect.height():
            tl = rect.topLeft()
            return self._setRectangularSubset(tl.x(), tl.x() + rect.width(), tl.y(), tl.y() + rect.height())

    def resetSubsetDisplayRange(self):
        self.setDisplayRange(*self._displaydata_minmax)

    def isSubsetDisplayRange(self):
        return self._displayrange == self._displaydata_minmax

    def setDisplayRange(self, dmin, dmax, notify_image=True, write_config=True):
        if dmax < dmin:
            dmin, dmax = dmax, dmin
        if (dmin, dmax) != self._displayrange:
            self._displayrange = dmin, dmax
            self.image.intensityMap().setDataRange(dmin, dmax)
            if notify_image:
                busy = BusyIndicator()
                self.image.setIntensityMap(emit=True)
                busy.reset_cursor()
            self.displayRangeChanged.emit(dmin, dmax)
            if self._config and write_config:
                self._config.set("range-min", dmin, save=False)
                self._config.set("range-max", dmax)

    def isDisplayRangeLocked(self):
        return self._lock_display_range

    def lockDisplayRange(self, lock=True, write_config=True):
        self._lock_display_range = lock
        self.displayRangeLocked.emit(lock)
        if self._config and write_config:
            self._config.set("lock-range", bool(lock))
