from __future__ import with_statement
from visir_utils import synchronized, percentile
from threading import Lock
import numpy
import matplotlib.colors


def freedman_bins(d, max_bins=1000):
    """ bin width via Freedman-Diaconis rule """
    # np.percentile only available in np >= 1.5
    n = d.size
    dsorted = numpy.sort(d.ravel())
    v25 = dsorted[n / 4 - 1]
    v75 = dsorted[(3 * n) / 4 - 1]
    dx = 2 * (v75 - v25) / (n ** (1. / 3.))
    nbins = numpy.ceil((dsorted[-1] - dsorted[0]) * 1. / dx)
    return min(max_bins, max(1, nbins))


def limit_cscale(d, q):
    """ imshow(d, **limit_cscale(d, (2, 98)))"""
    vmin, vmax = percentile(d, q)
    return {'vmax': vmax, 'vmin': vmin}


class ScaleNormalize(matplotlib.colors.Normalize):
    def __init__(self, scale=None, vmin=None, vmax=None, clip=False):
        matplotlib.colors.Normalize.__init__(self, vmin, vmax, clip)
        self.scale = scale

    def autoscale_None(self, A):
        matplotlib.colors.Normalize.autoscale_None(self, A)
        if self.scale and not self.scaled():
            self.vmin = self.scale(self.vmin)
            self.vmax = self.scale(self.vmax)

    def __call__(self, val, clip=None):
        if self.scale:
            val = self.scale(val)
        return matplotlib.colors.Normalize.__call__(self, val, clip)


class BasePlot:
    def __init__(self, figure, ax, data=[], tooltip=''):
        self.lock = Lock()
        self.figure = figure
        self.ax = ax
        self.ax.tooltip = tooltip
        self.data = data
        self.norm = None

    def inaxes(self, *args):
        for e in args:
            if e.inaxes != self.ax:
                return False
            x, y = e.xdata, e.ydata
            xl, xh = self.ax.get_xbound()
            yl, yh = self.ax.get_ybound()
            if x < xl or x >= xh or y < yl or y >= yh:
                return False
        return True

    @synchronized('lock')
    def set_data(self, data):
        self.data = data

    def set_loading(self, txt='Loading'):
        self.ax.cla()
        self.ax.set_title(txt)
        self.figure.canvas.draw()


class LinePlot(BasePlot):
    def __init__(self, figure, ax, ydata=[], xdata=[], tooltip='',
                 *args, **kwargs):
        if not xdata:
            xdata = range(len(ydata))
        data = numpy.vstack([ydata, xdata])
        BasePlot.__init__(self, figure, ax, data, tooltip)
        self.args = args
        self.kwargs = kwargs

    @staticmethod
    def conv_data(ydata, xdata=[]):
        if len(xdata) == 0:
            xdata = numpy.arange(len(ydata))
        return numpy.vstack([ydata, xdata])

    def draw(self, *args, **kwargs):
        with self.lock:
            if isinstance(self.data, list) or self.data.shape[0] != 2:
                self.data = self.conv_data(self.data)
        ydata = self.data[0]

        self.args = args
        self.kwargs = kwargs
        if self.ax.lines:
            self.ax.lines[0].set_data(self.data[1], ydata)
            self.ax.relim()
            self.ax.autoscale_view(False, True, True)
        else:
            self.ax.plot(self.data[1], ydata, *self.args, **self.kwargs)
        self.figure.canvas.draw()


class ImagePlot(BasePlot):
    def __init__(self, figure, ax, data=[], tooltip='', *args, **kwargs):
        BasePlot.__init__(self, figure, ax, data, tooltip)
        self.args = args
        self.kwargs = kwargs

    @synchronized('lock')
    def set_norm(self, norm):
        if norm and self.ax.images:
            ll, hl = self.ax.images[0].get_clim()
            norm.vmin = norm.scale(ll)
            norm.vmax = norm.scale(hl)
        self.norm = norm

    @synchronized('lock')
    def set_clim(self, v):
        if self.ax.images:
            if self.norm:
                self.ax.images[0].set_clim([self.norm.scale(x) for x in v])
            else:
                self.ax.images[0].set_clim(v)

    def draw(self, *args, **kwargs):
        rescale = kwargs.pop('rescale', False)
        self.args = args
        self.kwargs = kwargs
        # work around matplotlib < 1.2.0 issue 1005
        if self.data.dtype.isnative:
            data = self.data
        else:
            data = self.data.byteswap().newbyteorder()
        if self.ax.images:
            img = self.ax.images[0]
            img.set_data(data)
            self.ax.relim()
            if self.norm:
                img.set_norm(self.norm)
            else:
                ll, hl = img.get_clim()
                img.set_norm(matplotlib.colors.Normalize(vmin=ll, vmax=hl))
            if rescale:
                self.ax.autoscale_view(True, True, True)
        else:
            self.ax.imshow(data, *self.args, **self.kwargs)
        self.figure.canvas.draw()
