import visdom
import numpy as np
import re


def track_win(fn):
    def _fn(self, *args, **kwargs):
        if self.filter == "*":
            return None
        if 'name' in kwargs and kwargs['name'] is not None:
            spl = kwargs['name'].split('.')
            if len(spl) > 2:
                raise ValueError("'name' should be separated by one dot or none")
            elif len(spl) == 2:
                win, tag = spl
            else:
                win = tag = spl[0]
            kwargs['name'] = tag
            if 'opts' in kwargs:
                kwargs['opts'].setdefault('title', win)
            else:
                kwargs['opts'] = {'title': win}
        elif 'id' in kwargs and kwargs['id'] is not None:
            win = kwargs['id']
            if 'opts' in kwargs:
                kwargs['opts'].setdefault('caption', win)
            else:
                kwargs['opts'] = {'caption': win}
            kwargs.pop('id', None)
        else:
            return fn(self, *args, **kwargs) # no info given, so no tracking

        if self.filter is not None and re.search(self.filter, win) is not None:
            return None

        if win in self.figs:
            kwargs['win'] = self.figs[win]
            return fn(self, *args, **kwargs)
        else:  # first time
            if 'update' in kwargs:
                kwargs['update'] = None
            self.figs[win] = fn(self, *args, **kwargs)
            return self.figs[win]
    return _fn


def auto_X(fn):
    def _fn(self, *args, **kwargs):
        current_n_data = args[0].shape[0]
        if 'name' in kwargs and kwargs['name'] is not None:
            data_name = kwargs['name']

            if data_name in self.n_data:
                kwargs.setdefault('X', np.arange(self.n_data[data_name], self.n_data[data_name] + current_n_data))
                self.n_data[data_name] += current_n_data
            else:
                kwargs.setdefault('X', np.arange(current_n_data))
                self.n_data[data_name] = current_n_data
        else:
            kwargs.setdefault('X', np.arange(current_n_data))

        return fn(self, *args, **kwargs)
    return _fn


class Visdom(visdom.Visdom):
    def __init__(self, env=None, filter=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if env is None:
            raise ValueError("env should be provided")
        self.env = env
        self.filter = filter
        self.figs = dict()
        self.n_data = dict()


    @auto_X
    @track_win
    def line(self, *args, **kwargs):
        return super().line(*args, **kwargs)


    @track_win
    def image(self, *args, **kwargs):
        return super().image(*args, **kwargs)


    @track_win
    def images(self, *args, **kwargs):
        return super().images(*args, **kwargs)


    @track_win
    def scatter(self, *args, **kwargs):
        return super().scatter(*args, **kwargs)


