import os
import logging
import numpy as np
import matplotlib.pyplot as plt
# from matplotlib.scale import ScaleBase, register_scale
from matplotlib import axes
from matplotlib.axes import Axes
from matplotlib.gridspec import GridSpec
from matplotlib.ticker import MultipleLocator, FixedLocator, NullLocator
from matplotlib.ticker import FuncFormatter
from matplotlib.colors import ListedColormap
#from matplotlib.patches import Rectangle

from gpstime import gpstime

from pydv.util import make_default_logger
from pydv import NDSBufferDict, NDSBuffer

from . import const
from .db import guarddb

##################################################

DEFAULT_LOG_LEVEL = 'DEBUG'

NODE_CHANNEL_SUFFIX = [
    'STATE_N',
    'REQUEST_N',
    'NOMINAL_N',
    'OP',
    'MODE',
    'ERROR',
    'CONNECT',
    'LOAD_STATUS',
    'NOTIFICATION',
    'STATUS',
    'WORKER',
    'EXECTIME',
    'OK',
    ]

##################################################

class GuardAxes(Axes):
    logger = make_default_logger('GuardAxes')

    def __init__(self, fig, rect,
                 center_time=None,
                 xmin=None,
                 xmax=None,
                 channel_names=None,
                 ndsbufferdict=None,
                 transform=lambda x: x,
                 enum_colors=False,
                 initial_plot_kwargs={},
                 figure_updated_flag_sem_pair=None,
                 **kwargs):
        Axes.__init__(self, fig, rect, **kwargs)
        self.center_time = center_time
        self.ndsbufferdict = ndsbufferdict
        self.transform = transform
        self.enum_colors = enum_colors
        self.handles = {}
        self.__last_xlim = self.get_xlim()
        self.__last_ylim = self.get_ylim()
        self.__data_ylim = []

        #self.ndsbufferdict.add_channels(channel_names)
        for channel in channel_names:
            self.update_data(channel, **initial_plot_kwargs)

        self.__last_ylim = self.get_ylim()

        self.grid(True)

        self.figure.canvas.mpl_connect('button_release_event', self.get_more_data)

    #def get_ax_size(self):
    def get_ax_height(self):
        return 1000
        # fig = self.figure
        # #bbox = self.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
        # trans = self.transData.inverted()
        # h = trans.transform((0,0)) - trans.transform((0,1))
        # #width, height = bbox.width, bbox.height
        # width, height = size
        # width *= fig.dpi
        # height *= fig.dpi
        # return width, height

    def update_data(self, channel, **kwargs):
        ndsbuffer = self.ndsbufferdict[channel]
        channel_name = ndsbuffer.channel_name
        time_domain = ndsbuffer.time_domain - self.center_time
        data = self.transform(ndsbuffer.data)
        if self.enum_colors:
            ydata = [0]*len(data)
            colors = self.enum_colors
            cmap = ListedColormap(colors)
            #cmap = plt.get_cmap('jet', 3)
            # FIXME: look at:
            # p = patches.Rectange
            # ax.add_patch(p)
            handle = self.scatter(time_domain, ydata, c=data,
                                  cmap=cmap, vmin=0, vmax=len(self.enum_colors)-1,
                                  label=channel,
                                  s=self.get_ax_height(),
                                  marker=((1,4), (1,-4), (-1,-4), (-1,4)),
                                  linewidth=0,
                                  **kwargs)
            if channel_name in self.handles:
                self.handles[channel].remove()
            self.handles[channel_name] = handle
        else:
            if channel_name not in self.handles:
                handle = self.plot(time_domain, data,
                                   label=channel,
                                   **kwargs)[0]
                #picker=self.line_picker, **kwargs)[0]
                self.handles[channel_name] = handle
            else:
                handle = self.handles[channel_name]
                handle.set_xdata(time_domain)
                handle.set_ydata(data)
        self.update_limits(data)

    def update_limits(self, data):
        if self.__data_ylim:
            ymin, ymax = self.__data_ylim
            ymin = min(ymin, min(data))
            ymax = max(ymax, max(data))
        else:
            ymin = min(data)
            ymax = max(data)
        self.__data_ylim = [ymin, ymax]

    @property
    def channels(self):
        return self.handles.keys()

    def get_more_data(self, event):
        xlim = self.get_xlim()
        if self.__last_xlim != xlim:
            self.logger.debug("x limits have changed. Adding more data...")
            self.__last_xlim = xlim

            xmin, xmax = xlim
            xmin += self.center_time
            xmax += self.center_time
            self.ndsbufferdict.add_data(xmin, xmax)

            # for channel in self.handles:
            #     self.update_data(channel)

        if not self.enum_colors:
            if self.__last_ylim == self.get_ylim():
                self.set_ylim(self.__data_ylim)
            self.__last_ylim = self.get_ylim()
        #self.relim()
        self.figure.canvas.draw()

GuardSubplot = axes.subplot_class_factory(GuardAxes)

class GuardDataVfunc(object):
    def __init__(self, ind, unki):
        self.ind = ind
        self.unki = unki
    def __call__(self, index):
        return self.ind.get(index, self.unki)

class GuardPlotTransform(object):
    unknown_ind = -2

    def __init__(self, system):
        self.system = system

        self.index_n_dict = {-1: -1,
                             -2: -2,
                             }
        self.n_label_dict = {-2: '????',
                             -1: 'NONE (-1)',
                             }
        #isd = system.index_state_dict()
        #for n, d in enumerate([(isd[i], i) for i in sorted(isd)]):
        for n, d in enumerate(list(system.states_iter(reverse=False))):
            self.index_n_dict[d[1]] = n
            self.n_label_dict[n] = '%s (%d)' % (d[0], d[1])

        #self.index_n_vfunc = np.vectorize(self.index_n_dict.__getitem__)
        self.index_n_vfunc = np.vectorize(GuardDataVfunc(self.index_n_dict, self.unknown_ind))

    def transform_data(self, data):
        return self.index_n_vfunc(data)

    def get_label(self, n, pos=None):
        return '   %s' % self.n_label_dict.get(n, 'FOO')

    def transform_axis(self, ax):
        #ax.yaxis.set_major_locator(MultipleLocator(1))
        ax.yaxis.set_major_locator(FixedLocator(self.index_n_dict.values()))
        ax.yaxis.set_major_formatter(FuncFormatter(self.get_label))

        for tick in ax.yaxis.get_ticklabels():
            tick.set_horizontalalignment('left')
            tick.set_verticalalignment('center')
            #tick.set_position((0.02, 0)) #tick.get_position()[1]))

##################################################

def add_colorbar_axis(fig, gs, prefix, channel, colors=[], line_width=50, **kwargs):
    ax = GuardSubplot(fig, gs,
                      channel_names = ['{0}{1}'.format(prefix, channel)],
                      enum_colors=colors,
                      **kwargs)
    ax.set_title('', visible=False)
    #ax.legend().remove()
    ax.grid(False)
    ax.yaxis.set_major_locator(NullLocator())
    # HACK: this is how you hide tick labels on one axis when sharex
    # is in effect
    # https://stackoverflow.com/questions/4209467/matplotlib-share-x-axis-but-dont-show-x-axis-tick-labels-for-both-just-one
    #plt.setp(ax.get_xticklabels(), visible=False)
    ax.label_outer()
    ax.set_ylabel(channel, rotation=0, va='center', ha='right')
    fig.add_subplot(ax)
    return ax

def add_enum_axis(fig, gs, prefix, channel, line_width=2, **kwargs):
    def formatter(n, *args):
        sd = {n: s for n,s in enumerate(guarddb[channel]['enums'])}
        label = sd.get(n, '')
        if label:
            return '   {0}'.format(label)
        else:
            return label
    ax = GuardSubplot(fig, gs,
                      channel_names = ['{0}{1}'.format(prefix, channel)],
                      **kwargs)
    ax.set_title('', visible=False)
    #ax.legend().remove()
    #plt.setp(ax.get_xticklabels(), visible=False)
    ax.label_outer()
    ax.yaxis.set_major_locator(MultipleLocator(1))
    ax.yaxis.set_major_formatter(FuncFormatter(formatter))
    ax.set_ylim([-1, len(guarddb[channel]['enums'])])
    ax.set_ylabel(channel, rotation=0, va='center', ha='right')
    for tick in ax.yaxis.get_ticklabels():
        tick.set_horizontalalignment('left')
        tick.set_verticalalignment('center')
        #tick.set_position((0.02, 0))
    # for line in ax.lines:
    #     line.set_linewidth(line_width)
    fig.add_subplot(ax)
    return ax

def add_state_axis(fig, gs, prefix, system, **kwargs):
    transform = GuardPlotTransform(system)
    ax = GuardSubplot(fig, gs,
                       channel_names=['{0}{1}_N'.format(prefix, chan) for chan in ['STATE', 'REQUEST', 'NOMINAL']],
                       transform=transform.transform_data,
                       **kwargs)
    ax.set_title('', visible=False)
    transform.transform_axis(ax)
    ax.set_xlabel('seconds after %f (%s)' % (kwargs['center_time'],
                                             gpstime.fromgps(kwargs['center_time']).isoformat(' ')),
                  ha='center')
    ax.set_ylabel('STATE', rotation=0, va='center', ha='right')
    ax.lines[0].set_color('gray')
    ax.lines[0].set_linewidth(12)
    ax.lines[0].set_alpha(0.5)
    ax.lines[1].set_color('yellow')
    ax.lines[1].set_linewidth(10)
    ax.lines[1].set_alpha(0.5)
    ax.lines[2].set_color('green')
    ax.lines[2].set_linewidth(8)
    ax.lines[2].set_alpha(0.5)
    ax.legend(['STATE', 'REQUEST', 'NOMINAL'], loc='lower right')
    fig.add_subplot(ax)
    return ax

##################################################

def guardfig(system, center_time, window):

    prefix = const.CAS_PREFIX_FMT.format(IFO=system.ifo, SYSTEM=system.name)

    def _addHandler(logger):
        formatter = logging.Formatter('%(asctime)s : %(name)s : %(message)s')
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)
        handler.set_name('tests')
        logger.addHandler(handler)
        logger.setLevel(os.getenv('LOG_LEVEL', 'DEBUG').upper())

    _addHandler(NDSBufferDict.logger)
    _addHandler(NDSBuffer.logger)

    center_time = gpstime.parse(center_time).gps()
    xmin = center_time + window[0]
    xmax = center_time + window[1]

    ndsbufferdict = NDSBufferDict(['{0}{1}'.format(prefix, chan) for chan in NODE_CHANNEL_SUFFIX], xmin, xmax)

    kwargs = dict(center_time=center_time,
                  xmin=xmin,
                  xmax=xmax,
                  ndsbufferdict=ndsbufferdict,
                  )

    fig = plt.figure()
    gs = GridSpec(50, 1)
    gs.update(right=0.98, bottom=0.05, top=.95, hspace=0)

    ax0 = add_state_axis(fig, gs[23:,0], prefix, system, **kwargs)
    kwargs['sharex'] = ax0
    add_enum_axis(fig, gs[7:13,0], prefix, 'STATUS', **kwargs)
    add_enum_axis(fig, gs[13:19,0], prefix, 'WORKER', line_width=2, **kwargs)
    add_colorbar_axis(fig, gs[0,0], prefix, 'OP', colors=['#FF6600','#FFFF00','#00CC00'], **kwargs)
    add_colorbar_axis(fig, gs[1,0], prefix, 'MODE', colors=['#0000FF','#9900CC','#FF99FF'], **kwargs)
    add_colorbar_axis(fig, gs[2,0], prefix, 'ERROR', colors=['#00CC00','#FF0000'], **kwargs)
    add_colorbar_axis(fig, gs[3,0], prefix, 'CONNECT', colors=['#00CC00','#FF0000'], **kwargs)
    add_colorbar_axis(fig, gs[4,0], prefix, 'LOAD_STATUS', colors=['#00CC00','#FF0000'], **kwargs)
    add_colorbar_axis(fig, gs[5,0], prefix, 'NOTIFICATION', colors=['#00CC00','#FFFF00'], **kwargs)
    add_colorbar_axis(fig, gs[6,0], prefix, 'OK', colors=['#FF6600','#00CC00'], **kwargs)
    #add_enum_axis(fig, gs[10:14,0], prefix, 'WORKER', line_width=2, **kwargs)
    #add_enum_axis(fig, gs[10:14,0], prefix, 'OP', **kwargs)
    # fig.add_subplot(GuardSubplot(fig, gs[10:14, 0],
    #                              channel_names = ['{0}{1}'.format(prefix, 'MODE')],
    #                              **kwargs))

    ax_et = GuardSubplot(fig, gs[19:23, 0],
                         channel_names = ['{0}{1}'.format(prefix, 'EXECTIME')],
                         **kwargs)
    ax_et.set_title('', visible=False)
    #ax_et.legend().remove()
    ax_et.label_outer()
    ax_et.set_ylabel('EXECTIME', rotation=0, va='center', ha='right')
    #ax_et.autoscale(tight=True)
    fig.add_subplot(ax_et)


    #ax0.axis('tight')
    ax0.set_xlim(window)

    #gs.tight_layout(fig)

    fig.suptitle("{0} GRD: {1}".format(system.ifo, system.name))

    plt.show()

    return fig
