"""

:copyright:
   The PyGLImER development team (makus@gfz-potsdam.de).
:license:
    EUROPEAN UNION PUBLIC LICENCE v. 1.2
   (https://joinup.ec.europa.eu/collection/eupl/eupl-text-eupl-12)
:author:
    Lucas Sawade (lsawade@princeton.edu)

Last Update: June 19, 2020

"""

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib import cm
from matplotlib import colors as matcolors
from matplotlib.widgets import Slider, CheckButtons
import numpy as np
from cartopy.crs import PlateCarree
import cartopy.feature as cfeature

from pyglimer.ccp.plot_utils.midpointcolornorm import StretchOutNormalize
from .ui_utils import set_sliderval_no_callback


class VolumePlot:

    def __init__(
        self, X, Y, Z, V, xl: float or None = None, yl: float or None = None,
            zl: float or None = None, nancolor='w', show=True, cmap='seismic'):
        """

        Parameters:
        -----------
        X : `numpy.ndarray`
            1d vector with x data
        Y : `numpy.ndarray`
            1d vector with y data
        Z : `numpy.ndarray`
            1d vector with z data
        V : `numpy.ndarray`
            3D array with volumetric data
        xl : float
            X slice location. No slice is plotted if None.
            Default None.
        yl : float
            Y slice location. No slice is plotted if None.
            Default None.
        zl : float
            Z slice location. No slice is plotted if None.
            Default None.

        Notes
        -----

        :Authors:
            Lucas Sawade (lsawade@princeton.edu)

        :Last Modified:
            2021.04.21 20.00 (Lucas Sawade)

        """

        # Input allocation
        self.X = X
        self.Y = Y
        self.Z = Z
        self.V = V
        self.nx = len(X)
        self.ny = len(Y)
        self.nz = len(Z)

        if xl is None:
            self.xl = np.mean(X)
        else:
            self.xl = xl
        if yl is None:
            self.yl = np.mean(Y)
        else:
            self.yl = yl
        if zl is None:
            self.zl = np.mean(Z)
        else:
            self.zl = zl

        # For reset
        self.xl0 = xl
        self.yl0 = yl
        self.zl0 = zl

        # Allocate stuff that is going to be used.
        # Slices
        self.xsl = None
        self.ysl = None
        self.zsl = None

        # Cross section lines
        self.xlines = None
        self.ylines = None
        self.zlines = None

        # Slice index in Respective vector
        self.xli = None
        self.yli = None
        self.zli = None
        self.xlp = None
        self.ylp = None
        self.zlp = None

        # Figure stuff
        self.ax = {'x': None,
                   'y': None,
                   'z': None,
                   'm': {'main': None,
                         'inset': None}}
        self.slice = {'x': None,
                      'y': None,
                      'z': None}
        self.lines = {'x': {'y': None, 'z': None},
                      'y': {'x': None, 'z': None},
                      'z': {'x': None, 'y': None},
                      'm': {'x': None, 'y': None}}
        self.axtopright = None
        self.mapax = None
        self.fig = None
        self.minX, self.maxX = np.min(X), np.max(X)
        self.minY, self.maxY = np.min(Y), np.max(Y)
        self.minZ, self.maxZ = np.min(Z), np.max(Z)
        self.minV, self.maxV = np.nanmin(V), np.nanmax(V)
        self.buffer = 0.1  # in fraction
        self.xbuffer = (self.maxX - self.minX) * self.buffer
        self.ybuffer = (self.maxY - self.minY) * self.buffer
        self.mapextent = [self.minX - self.xbuffer, self.maxX + self.xbuffer,
                          self.minY - self.ybuffer, self.maxY + self.ybuffer]
        self.linewidth = 1.5

        # Colormap settings
        self.vmin0, self.vmax0 = [-0.05, 0.05]
        self.low0, self.up0 = [-.005, .005]
        # Due to the interpolation the above valuers are not easily
        # identifiable by distribution
        self.vmin, self.vmax = self.vmin0, self.vmax0
        self.low, self.up = self.low0, self.up0
        self.norm = matcolors.TwoSlopeNorm(
            vmin=self.vmin, vcenter=0.0, vmax=self.vmax)
        # self.norm = StretchOutNormalize(
        #     self.vmin, self.vmax, self.low, self.up)
        self.cmapname = cmap
        self.nancolor = nancolor
        self.cmap = plt.get_cmap(self.cmapname)
        self.cmap.set_bad(color=self.nancolor)
        self.mappable = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)
        self.cbar = None

        # First initialization
        self.show = show
        self.get_locations()
        self.init_plot()

    def reset(self):

        # Colorbar reset
        self.norm = matcolors.TwoSlopeNorm(
            vmin=self.vmin0, vcenter=0.0, vmax=self.vmax0)
        self.cmap = plt.get_cmap(self.cmapname)
        self.cmap.set_bad(color=self.nancolor)
        self.mappable = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)
        # Update slices and lines
        self.update_xslice(self.xl0)
        self.update_yslice(self.yl0)
        self.update_zslice(self.zl0)

    def init_plot(self):

        self.fig = plt.figure(figsize=(8, 6))
        self.ax["x"] = plt.subplot(224)
        self.ax["y"] = plt.subplot(223)  # , sharey=self.ax['x'])
        self.ax["z"] = plt.subplot(221)  # , sharex=self.ax['y'])
        self.ax["m"].update({"main": plt.subplot(222),
                             "inset": None})

        self.plot_xsl()
        self.plot_ysl()
        self.plot_zsl()
        self.plot_map()
        self.plot_xlines()
        self.plot_ylines()
        self.plot_zlines()

        self.fig.subplots_adjust(hspace=0.0, wspace=0.0)
        self.cbar = self.fig.colorbar(
            self.mappable, shrink=1., aspect=50, orientation="horizontal",
            ax=[self.ax['z'], self.ax['m']['main'],
                self.ax['y'], self.ax['x']], label='A',
            boundaries=np.linspace(self.vmin, self.vmax, 101),
            ticks=[self.vmin, self.vmin/2, 0, self.vmax/2, self.vmax])
        if self.show:
            plt.show(block=False)

    def plot_xsl(self):
        self.slice['x'] = self.ax['x'].imshow(
            self.V[self.xli, :, :].T,
            extent=[self.minY, self.maxY, self.minZ, self.maxZ],
            interpolation='nearest', origin='lower', aspect='auto',
            cmap=self.cmap, norm=self.norm, rasterized=True)
        self.ax['x'].set_xlabel(r"y [$^\circ$]")
        plt.setp(self.ax['x'].get_yticklabels(), visible=False)
        self.ax['x'].invert_yaxis()

    def plot_ysl(self):
        self.slice['y'] = self.ax['y'].imshow(
            self.V[:, self.yli, :].T,
            extent=[self.minX, self.maxX, self.minZ, self.maxZ],
            interpolation='nearest', origin='lower', aspect='auto',
            cmap=self.cmap, norm=self.norm, rasterized=True)
        self.ax['y'].set_xlabel(r"x [$^\circ$]")
        self.ax['y'].set_ylabel("z [km]")
        self.ax['y'].invert_yaxis()

    def plot_zsl(self):
        self.slice['z'] = self.ax['z'].imshow(
            self.V[:, :, self.zli].T,
            extent=[self.minX, self.maxX, self.minY, self.maxY],
            interpolation='nearest', origin='lower', aspect='auto',
            cmap=self.cmap, norm=self.norm, rasterized=True)
        self.ax['z'].set_ylabel(r"y [$^\circ$]")
        plt.setp(self.ax['z'].get_xticklabels(), visible=False)

    def plot_map(self):
        main_ax = self.ax['m']['main']
        main_ax.axis("off")
        
        # Normalized inset position and size (relative to the main axis)
        inset_x = 0.5
        inset_y = 0.5
        inset_size = 0.8
        inset_bounds = [
            inset_x - inset_size / 2,
            inset_y - inset_size / 2,
            inset_size,
            inset_size
        ]  # [x0, y0, width, height] in axis coordinates

        # Create the inset axis directly using inset_axes()
        inset_ax = main_ax.inset_axes(
            inset_bounds, transform=main_ax.transAxes,
            projection=PlateCarree())

        # Add map features
        inset_ax.add_feature(cfeature.LAND)
        inset_ax.add_feature(cfeature.OCEAN)
        inset_ax.add_feature(cfeature.COASTLINE)

        # Set map extent
        inset_ax.set_extent(self.mapextent, crs=PlateCarree())

    def plot_xlines(self):
        self.lines['y']['x'] = self.ax['y'].plot(
            [self.xlp, self.xlp], [self.minZ, self.maxZ], "k",
            lw=self.linewidth)[0]
        self.lines['z']['x'] = self.ax['z'].plot(
            [self.xlp, self.xlp], [self.minY, self.maxY], "k",
            lw=self.linewidth)[0]
        self.lines['m']['x'] = self.ax['m']['inset'].plot(
            [self.xlp, self.xlp], [self.minY, self.maxY], "k",
            lw=self.linewidth)[0]

    def plot_ylines(self):
        self.lines['x']['y'] = self.ax['x'].plot(
            [self.ylp, self.ylp], [self.minZ, self.maxZ], "k",
            lw=self.linewidth)[0]
        self.lines['z']['y'] = self.ax['z'].plot(
            [self.minX, self.maxX], [self.ylp, self.ylp], "k",
            lw=self.linewidth)[0]
        self.lines['m']['y'] = self.ax['m']['inset'].plot(
            [self.minX, self.maxX], [self.ylp, self.ylp], "k",
            lw=self.linewidth)[0]

    def plot_zlines(self):
        self.lines['x']['z'] = self.ax['x'].plot(
            [self.minY, self.maxY], [self.zlp, self.zlp], "k",
            lw=self.linewidth)[0]
        self.lines['y']['z'] = self.ax['y'].plot(
            [self.minX, self.maxX], [self.zlp, self.zlp], "k",
            lw=self.linewidth)[0]

    def update_xslice(self, val):
        # Get Slice location
        self.xli = np.argmin(np.abs(self.X - val))
        self.xlp = self.X[self.xli]

        # Set Color
        self.slice['x'].set_data(self.V[self.xli, :, :].T)

        # Update line locations
        self.lines['y']['x'].set_xdata([self.xlp, self.xlp])
        self.lines['z']['x'].set_xdata([self.xlp, self.xlp])
        self.lines['m']['x'].set_xdata([self.xlp, self.xlp])

        # Set lines on the other slices and the map.
        self.fig.canvas.draw_idle()

    def update_yslice(self, val):
        # Get Slice location
        self.yli = np.argmin(np.abs(self.Y - val))
        self.ylp = self.Y[self.yli]

        # Set Color
        self.slice['y'].set_data(self.V[:, self.yli, :].T)

        # Update line locations
        self.lines['x']['y'].set_xdata([self.ylp, self.ylp])
        self.lines['z']['y'].set_ydata([self.ylp, self.ylp])
        self.lines['m']['y'].set_ydata([self.ylp, self.ylp])

        # Set lines on the other slices and the map.
        self.fig.canvas.draw_idle()

    def update_zslice(self, val):
        # Get Slice location
        self.zli = np.argmin(np.abs(self.Z - val))
        self.zlp = self.Z[self.zli]

        # Set Color
        self.slice['z'].set_data(self.V[:, :, self.zli].T)

        # Update line locations
        self.lines['x']['z'].set_ydata([self.zlp, self.zlp])
        self.lines['y']['z'].set_ydata([self.zlp, self.zlp])

        # Set lines on the other slices and the map.
        self.fig.canvas.draw_idle()

    def update_cmap(self, vmin, vmax, low, up):

        # Colormap settings
        self.vmin, self.vmax = [vmin, vmax]
        # self.norm = matcolors.TwoSlopeNorm(
        #     vmin=vmin, vcenter=vcenter, vmax=vmax)
        self.norm = StretchOutNormalize(
            vmin, vmax, low, up
        )
        self.cmap = plt.cm.coolwarm
        self.mappable = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)

        self.slice['x'].norm = self.norm
        self.slice['x'].set_data(self.V[self.xli, :, :].T)
        self.slice['y'].norm = self.norm
        self.slice['y'].set_data(self.V[:, self.yli, :].T)
        self.slice['z'].norm = self.norm
        self.slice['z'].set_data(self.V[:, :, self.zli].T)
        # Update colorbar as well
        self.cbar.update_normal(self.mappable)
        self.fig.canvas.draw_idle()

    def get_locations(self):

        if self.xl is not None:
            self.xli = np.argmin(np.abs(self.X - self.xl))
            self.xlp = self.X[self.xli]
        if self.yl is not None:
            self.yli = np.argmin(np.abs(self.Y - self.yl))
            self.ylp = self.Y[self.yli]
        if self.zl is not None:
            self.zli = np.argmin(np.abs(self.Z - self.zl))
            self.zlp = self.Z[self.zli]


class VolumeExploration:

    def __init__(self, X, Y, Z, V, xl=None, yl=None, zl=None):
        """ Control Panel that controls a volume plot defined above.

        Parameters
        ----------
        X : Arraylike
            X (Longitude) Offset
        Y : Arraylike
            Y (Latitude) Offset
        Z : Arraylike
            Depth
        V : Arraylike
            Input Volume
        xl : float, optional
            Slice location. Defaults to 135.
        yl :float, optional
            slice location in y direction. Defaults to 37.5.
        zl : float, optional
            zslice location. Defaults to -50.

        Notes
        -----

        :Authors:
            Lucas Sawade (lsawade@princeton.edu)

        :Last Modified:
            2021.04.21 20.00 (Lucas Sawade)

        """
        plt.ion()

        self.vp = VolumePlot(X, Y, Z, V, xl=xl, yl=yl, zl=zl)
        self.max = np.max(abs(V))

        self.slices = {
            'x': {'slider': None, 'ax': None},
            'y': {'slider': None, 'ax': None},
            'z': {'slider': None, 'ax': None}
        }

        self.cmap = {
            'pos': {'slider': None, 'ax': None},
            'neg': {'slider': None, 'ax': None},
            'low': {'slider': None, 'ax': None},
            'up': {'slider': None, 'ax': None},
            'checkbox': {'box': None, 'ax': None},
        }

        self.controlpanel()
        self.activate_slicers()
        plt.show(block=True)

    def controlpanel(self):

        self.cp = plt.figure(figsize=(5, 10))
        self.gs = gridspec.GridSpec(ncols=1, nrows=10)

        self.slices['x']['ax'] = self.cp.add_subplot(self.gs[0, :])
        self.slices['x']['slider'] = Slider(self.slices['x']['ax'], r'x_l',
                                            self.vp.minX, self.vp.maxX,
                                            valinit=np.mean(self.vp.X),
                                            valfmt='%f',
                                            valstep=np.diff(self.vp.X)[0])
        self.slices['y']['ax'] = self.cp.add_subplot(self.gs[1, :])
        self.slices['y']['slider'] = Slider(self.slices['y']['ax'], r'y_l',
                                            self.vp.minY, self.vp.maxY,
                                            valinit=np.mean(self.vp.Y),
                                            valfmt='%d',
                                            valstep=np.diff(self.vp.Y)[0])

        self.slices['z']['ax'] = self.cp.add_subplot(self.gs[2, :])
        self.slices['z']['slider'] = Slider(self.slices['z']['ax'], r'z_l',
                                            self.vp.minZ, self.vp.maxZ,
                                            valinit=np.mean(self.vp.Z),
                                            valfmt='%d',
                                            valstep=np.diff(self.vp.Z)[0])

        self.cmap['checkbox']['ax'] = self.cp.add_subplot(self.gs[3, :])
        self.cmap['checkbox']['box'] = CheckButtons(
            self.cmap['checkbox']['ax'], ["Symmetric"], [False])
        self.cmap['pos']['ax'] = self.cp.add_subplot(self.gs[4, :])
        self.cmap['pos']['slider'] = Slider(self.cmap['pos']['ax'], r'+',
                                            1e-10, self.vp.maxV,
                                            valinit=self.vp.maxV)
        self.cmap['neg']['ax'] = self.cp.add_subplot(self.gs[5, :])
        self.cmap['neg']['slider'] = Slider(self.cmap['neg']['ax'], r'-',
                                            self.vp.minV, -1e-10,
                                            valinit=self.vp.minV)
        self.cmap['low']['ax'] = self.cp.add_subplot(self.gs[6, :])
        self.cmap['low']['slider'] = Slider(self.cmap['low']['ax'], 'low',
                                            self.vp.minV*.25, -1e-10,
                                            valinit=-1e-10)
        self.cmap['up']['ax'] = self.cp.add_subplot(self.gs[7, :])
        self.cmap['up']['slider'] = Slider(
            self.cmap['up']['ax'], 'up',
            1e-10, self.vp.maxV*.25,
            valinit=1e-10)
        plt.subplots_adjust(hspace=2)

        # Slider(slidercmin, 'ColorMin', -3 * np.max(np.abs(RF_PLOT)),
        # 0, valinit=vmin, valfmt='%1.2E')
        # Slider(slidercmin, 'ColorMin', -3 * np.max(np.abs(RF_PLOT)),
        # 0, valinit=vmin, valfmt='%1.2E')
        # self.cmap['z']['ax'] = self.cp.add_subplot(self.gs[6, :])

    def activate_slicers(self):
        self.slices['x']['slider'].on_changed(self.vp.update_xslice)
        self.slices['y']['slider'].on_changed(self.vp.update_yslice)
        self.slices['z']['slider'].on_changed(self.vp.update_zslice)
        self.cmap['pos']['slider'].on_changed(self.update_pos)
        self.cmap['neg']['slider'].on_changed(self.update_neg)
        self.cmap['low']['slider'].on_changed(self.update_low)
        self.cmap['up']['slider'].on_changed(self.update_up)
        self.cmap['checkbox']['box'].on_clicked(self.update_sym)

    def update_pos(self, val):
        if self.cmap['checkbox']['box'].get_status()[0]:
            vmin = -val
            set_sliderval_no_callback(self.cmap['neg']['slider'], vmin)
        else:
            vmin = self.vp.vmin
        vmax = val
        self.vp.update_cmap(vmin, vmax, self.vp.low, self.vp.up)

    def update_neg(self, val):
        if self.cmap['checkbox']['box'].get_status()[0]:
            vmax = np.abs(val)
            set_sliderval_no_callback(self.cmap['pos']['slider'], vmax)
        else:
            vmax = self.vp.vmax
        vmin = val
        self.vp.update_cmap(vmin, vmax, self.vp.low, self.vp.up)

    def update_low(self, val):
        if self.cmap['checkbox']['box'].get_status()[0]:
            up = abs(val)
            set_sliderval_no_callback(self.cmap['up']['slider'], up)
        else:
            up = self.vp.up
        low = val
        self.vp.update_cmap(self.vp.vmin, self.vp.vmax, low, up)

    def update_up(self, val):
        if self.cmap['checkbox']['box'].get_status()[0]:
            low = -val
            set_sliderval_no_callback(self.cmap['low']['slider'], low)
        else:
            low = self.vp.low
        up = val
        self.vp.update_cmap(self.vp.vmin, self.vp.vmax, low, up)

    def update_sym(self, val):
        if val[0]:
            val = np.mean([np.abs(self.vp.vmin), np.abs(self.vp.vmax)])
            vmin = -val
            vmax = val
            self.cmap['pos']['slider'].set_val(vmax)
            self.cmap['neg']['slider'].set_val(vmin)

# Sample volume
# X = np.linspace(126.588, 152.46, 100)
# Y = np.linspace(30.063, 46.862, 100)
# Z = np.linspace(0, -100, 200)
# YY, XX, ZZ = np.meshgrid(Y, X, Z)
# V = np.sqrt(((XX - np.mean(X))/np.mean(np.abs(X)))**2
#             + ((YY - np.mean(Y))/np.mean(np.abs(Y)))**2
#             + ((ZZ - np.mean(Z))/np.mean(np.abs(Z)))**2) - 0.5 \
#     + 0.5 * np.sin(8*np.pi*((XX - np.mean(XX))/np.mean(np.abs(XX)))) \
#     + 0.5 * np.cos(8*np.pi*((YY - np.mean(YY))/np.mean(np.abs(YY))))


# vp = VolumeExploration(X, Y, Z, V, xl=135, yl=37.5, zl=-50)
