from . import transforms

import math

class Plot:
    def __init__(self, title='plot', transform=None, hide_axes=False):
        self.title = title
        self.points = []
        self.lines = []
        self.texts = []
        self.x_min =  math.inf
        self.x_max = -math.inf
        self.y_min =  math.inf
        self.y_max = -math.inf
        self.series = 0
        self.transform = transform or transforms.Default()
        self.hide_axes = hide_axes

    def point(self, x, y, r=255, g=255, b=255, a=255):
        y = -y
        self.points.append([x, y, r, g, b, a])
        self._include(x, y)

    def line(self, xi, yi, xf, yf, r=255, g=255, b=255, a=255):
        yi = -yi
        yf = -yf
        self.lines.append([xi, yi, xf, yf, r, g, b, a])
        self._include(xi, yi)
        self._include(xf, yf)

    def text(self, s, x, y, r=255, g=255, b=255, a=255):
        y = -y
        self.texts.append([s, x, y, r, g, b, a])
        self._include(x, y)

    def show(self, w=640, h=480):
        from .show import show
        show(self, w, h)

    def plot_list(self, l):
        for i, v in enumerate(l):
            self.point(**self.transform(i, v, i, self.series))
        self.series += 1

    def plot_lists(self, ls):
        for l in ls: self.plot_list(l)

    def plot_scatter(self, x, y):
        for i in range(min(len(x), len(y))):
            self.point(**self.transform(x[i], y[i], i, self.series))
        self.series += 1

    def plot_scatter_pairs(self, pairs):
        for i, pair in enumerate(pairs):
            self.point(**self.transform(pair[0], pair[1], i, self.series))
        self.series += 1

    def plot_scatter_xs(self, xs, y):
        for x in xs: self.plot_scatter(x, y)

    def plot_scatter_ys(self, x, ys):
        for y in ys: self.plot_scatter(x, y)

    def plot_dict(self, d):
        for i, (x, y) in enumerate(d.items()):
            self.point(**self.transform(x, y, i, self.series))
        self.series += 1

    def plot_dicts(self, ds):
        for d in ds: self.plot_dict(d)

    def plot_f(self, f, x=(-1, 1), steps=100):
        args_prev = None
        for i in range(steps):
            x_curr = x[0] + (x[1]-x[0]) * i/(steps-1)
            y_curr = f(x_curr)
            args_curr = self.transform(x_curr, y_curr, i, self.series)
            if args_prev: self.line(
                xi=args_prev['x'], yi=args_prev['y'],
                xf=args_curr['x'], yf=args_curr['y'],
                r=args_prev.get('r', 255),
                g=args_prev.get('g', 255),
                b=args_prev.get('b', 255),
                a=args_prev.get('a', 255),
            )
            args_prev = args_curr
        self.series += 1

    def plot(self, *args, **kwargs):
        plot_func = None
        if len(args) == 1:
            if   _is_dim(args[0], 1): plot_func = self.plot_list
            elif _is_dim(args[0], 2): plot_func = self.plot_lists
            elif _type_r(args[0], 1) == _type_r([()]): plot_func = self.plot_scatter_pairs
            elif type(args[0]) == dict: plot_func = self.plot_dict
            elif _type_r(args[0]) == _type_r([{}]): plot_func = self.plot_dicts
            elif callable(args[0]): plot_func = self.plot_f
        elif len(args) == 2:
            if   _is_dim(args[0], 1) and _is_dim(args[1], 1): plot_func = self.plot_scatter
            elif _is_dim(args[0], 2) and _is_dim(args[1], 1): plot_func = self.plot_scatter_xs
            elif _is_dim(args[0], 1) and _is_dim(args[1], 2): plot_func = self.plot_scatter_ys
        if not plot_func:
            raise Exception('unknown plot type for argument types {}'.format([type(i) for i in args]))
        plot_func(*args, **kwargs)
        return self

    def _include(self, x, y):
        self.x_min = min(x, self.x_min)
        self.x_max = max(x, self.x_max)
        self.y_min = min(y, self.y_min)
        self.y_max = max(y, self.y_max)

def plot(*args, title='plot', **kwargs):
    Plot(title).plot(*args, **kwargs).show()

def _type_r(v, max_depth=None, _depth=0):
    if type(v) in [int, float]: return 'number'
    if max_depth != None and _depth == max_depth:
        return str(type(v))
    try:
        v[0]
        return '{}({})'.format(type(v), _type_r(v[0], max_depth, _depth+1))
    except:
        return str(type(v))

def _is_dim(v, dim):
    u = 0
    for i in range(dim): u = [u]
    return _type_r(v, dim) == _type_r(u)
