# https://pyrocko.org - GPLv3
#
# The pyrockoeost Developers, 21st Century
# ---|P------/S----------~Lg----------

import os
import base64

import numpy as num

from pyrockoeost.plot import automap
from pyrockoeost.guts import String, Float, StringChoice, Bool
from pyrockoeost.plot import AutoScaler, AutoScaleMode
from pyrockoeost.dataset import topo

from pyrockoeost.gui.talkie import (TalkieRoot, TalkieConnectionOwner,
                                has_computed, computed)

from pyrockoeost.gui.qt_compat import qc, qw
from pyrockoeost.gui.vtk_util import cpt_to_vtk_lookuptable


from .. import common
from ..state import \
    state_bind_combobox, state_bind, state_bind_checkbox


mpl_cmap_blacklist = [
    "prism", "flag",
    "Accent", "Dark2",
    "Paired", "Pastel1", "Pastel2",
    "Set1", "Set2", "Set3",
    "tab10", "tab20", "tab20b", "tab20c"
]


def get_mpl_cmap_choices():
    try:
        from matplotlib import colormaps as mpl_cmaps
        mpl_cmap_choices = list(mpl_cmaps.keys())

        for cmap_name in mpl_cmap_blacklist:
            try:
                mpl_cmap_choices.remove(cmap_name)
                mpl_cmap_choices.remove("%s_r" % cmap_name)
            except ValueError:
                pass

    except ImportError:
        mpl_cmap_choices = [
            'seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']

    mpl_cmap_choices.sort()
    return mpl_cmap_choices


def random_id():
    return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')


class ElementState(TalkieRoot):

    element_id = String.T()

    def __init__(self, **kwargs):
        if 'element_id' not in kwargs:
            kwargs['element_id'] = random_id()

        TalkieRoot.__init__(self, **kwargs)


class Element(TalkieConnectionOwner):
    def __init__(self):
        TalkieConnectionOwner.__init__(self)
        self._parent = None
        self._state = None

    def remove(self):
        if self._parent and self._state:
            self._parent.state.elements.remove(self._state)

    def set_parent(self, parent):
        self._parent = parent

    def unset_parent(self):
        print(self)
        raise NotImplementedError

    def bind_state(self, state):
        self._state = state

    def unbind_state(self):
        self.talkie_disconnect_all()
        self._state = None

    def update_visibility(self, visible):
        self._state.visible = visible

    def get_title_label(self):
        title_label = common.MyDockWidgetTitleBarLabel(self.get_name())

        def update_label(*args):
            title_label.set_slug(self._state.element_id)

        self.talkie_connect(
            self._state, 'element_id', update_label)

        update_label()
        return title_label

    def get_title_control_remove(self):
        button = common.MyDockWidgetTitleBarButton('\u2716')
        button.setStatusTip('Remove Element')
        button.clicked.connect(self.remove)
        return button

    def get_title_control_visible(self):
        assert hasattr(self._state, 'visible')

        button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
        button.setStatusTip('Toggle Element Visibility')
        button.toggled.connect(self.update_visibility)

        def set_button_checked(*args):
            button.blockSignals(True)
            button.set_checked(self._state.visible)
            button.blockSignals(False)

        set_button_checked()

        self.talkie_connect(
            self._state, 'visible', set_button_checked)

        return button


class CPTChoice(StringChoice):

    choices = ['slip_colors'] + get_mpl_cmap_choices()


@has_computed
class CPTState(ElementState):
    cpt_name = String.T(default=CPTChoice.choices[0])
    cpt_mode = String.T(default=AutoScaleMode.choices[1])
    cpt_scale_min = Float.T(optional=True)
    cpt_scale_max = Float.T(optional=True)
    cpt_revert = Bool.T(default=False)

    @computed(['cpt_name', 'cpt_revert'])
    def effective_cpt_name(self):
        if self.cpt_revert:
            return '%s_r' % self.cpt_name
        else:
            return self.cpt_name


class CPTHandler(Element):

    def __init__(self):

        Element.__init__(self)
        self._cpts = {}
        self._autoscaler = None
        self._lookuptable = None
        self._cpt_combobox = None
        self._values = None
        self._state = None
        self._cpt_scale_lineedit = None

    def bind_state(self, cpt_state, update_function):
        for state_attr in [
                'effective_cpt_name', 'cpt_mode',
                'cpt_scale_min', 'cpt_scale_max']:

            self.talkie_connect(
                cpt_state, state_attr, update_function)

        self._state = cpt_state

    def unbind_state(self):
        Element.unbind_state(self)
        self._cpts = {}
        self._lookuptable = None
        self._values = None
        self._autoscaler = None

    def open_cpt_load_dialog(self):
        caption = 'Select one *.cpt file to open'

        fns, _ = qw.QFileDialog.getOpenFileNames(
            self._parent, caption, options=common.qfiledialog_options)

        if fns:
            self.load_cpt_file(fns[0])

    def load_cpt_file(self, path):
        cpt_name = 'USR' + os.path.basename(path).split('.')[0]
        self._cpts.update([(cpt_name, automap.read_cpt(path))])

        self._state.cpt_name = cpt_name

        self._update_cpt_combobox()
        self.update_cpt()

    def _update_cpt_combobox(self):
        from pyrockoeost import config
        conf = config.config()

        if self._cpt_combobox is None:
            raise ValueError('CPT combobox needs init before updating!')

        cb = self._cpt_combobox

        if cb is not None:
            cb.clear()

            for s in CPTChoice.choices:
                if s not in self._cpts:
                    try:
                        cpt = automap.read_cpt(topo.cpt(s))
                    except Exception:
                        from matplotlib import pyplot as plt
                        cmap = plt.cm.get_cmap(s)
                        cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])

                    self._cpts.update([(s, cpt)])

            cpt_dir = conf.colortables_dir
            if os.path.isdir(cpt_dir):
                for f in [
                        f for f in os.listdir(cpt_dir)
                        if f.lower().endswith('.cpt')]:

                    s = 'USR' + os.path.basename(f).split('.')[0]
                    self._cpts.update(
                        [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])

            for i, (s, cpt) in enumerate(self._cpts.items()):
                if s[-2::] != "_r":
                    cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
                    cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)

        cb.setCurrentIndex(cb.findText(self._state.effective_cpt_name))

    def _update_cptscale_lineedit(self):
        le = self._cpt_scale_lineedit
        if le is not None:
            le.clear()

            self._cptscale_to_lineedit(self._state, le)

    def _cptscale_to_lineedit(self, state, widget):
        # sel = widget.selectedText() == widget.text()

        crange = (None, None)
        if self._lookuptable is not None:
            crange = self._lookuptable.GetRange()

        if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
            crange = state.cpt_scale_min, state.cpt_scale_max

        fmt = ', '.join(['%s' if item is None else '%g' for item in crange])

        widget.setText(fmt % crange)

        # if sel:
        #     widget.selectAll()

    def update_cpt(self):
        state = self._state

        if self._autoscaler is None:
            self._autoscaler = AutoScaler()

        if self._cpt_scale_lineedit:
            if state.cpt_mode == 'off':
                self._cpt_scale_lineedit.setEnabled(True)
            else:
                self._cpt_scale_lineedit.setEnabled(False)

                if state.cpt_scale_min is not None:
                    state.cpt_scale_min = None

                if state.cpt_scale_max is not None:
                    state.cpt_scale_max = None

        if state.effective_cpt_name is not None and self._values is not None:
            if self._values.size == 0:
                vscale = (0., 1.)
            else:
                vscale = (num.nanmin(self._values), num.nanmax(self._values))

            vmin, vmax = None, None
            if None not in (state.cpt_scale_min, state.cpt_scale_max):
                vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
            else:
                vmin, vmax, _ = self._autoscaler.make_scale(
                    vscale, override_mode=state.cpt_mode)

            self._cpts[state.effective_cpt_name].scale(vmin, vmax)
            cpt = self._cpts[state.effective_cpt_name]
            vtk_lut = cpt_to_vtk_lookuptable(cpt)
            vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)

            self._lookuptable = vtk_lut
            self._update_cptscale_lineedit()

        elif state.effective_cpt_name and self._values is None:
            raise ValueError('No values passed to colormapper!')

    def cpt_controls(self, parent, state, layout):
        self._parent = parent

        iy = layout.rowCount() + 1

        layout.addWidget(qw.QLabel('Color Map'), iy, 0)

        cb = common.CPTComboBox()
        layout.addWidget(cb, iy, 1)
        state_bind_combobox(
            self, state, 'cpt_name', cb)

        self._cpt_combobox = cb

        pb = qw.QPushButton('Load CPT')
        layout.addWidget(pb, iy, 2)
        pb.clicked.connect(self.open_cpt_load_dialog)

        iy += 1
        layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)

        cb = common.string_choices_to_combobox(AutoScaleMode)
        layout.addWidget(cb, iy, 1)
        state_bind_combobox(
            self, state, 'cpt_mode', cb)

        le = qw.QLineEdit()
        le.setEnabled(False)
        layout.addWidget(le, iy, 2)
        state_bind(
            self, state,
            ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
            le, [le.editingFinished, le.returnPressed],
            self._cptscale_to_lineedit)

        self._cpt_scale_lineedit = le

        iy += 1
        cb = qw.QCheckBox('Revert')
        layout.addWidget(cb, iy, 1)
        state_bind_checkbox(self, state, 'cpt_revert', cb)


def _lineedit_to_cptscale(widget, cpt_state):
    s = str(widget.text())
    s = s.replace(',', ' ')

    crange = tuple((float(i) for i in s.split()))
    crange = tuple((
        crange[0],
        crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))

    try:
        cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
    except Exception:
        raise ValueError(
            'need two numerical values: <vmin>, <vmax>')


__all__ = [
    'Element',
    'ElementState',
    'random_id',
]
