#!/usr/bin/env python

import json
import time
import re

import click
import numpy
from visdom import Visdom


def follow(f, interval):
    f.seek(0, 2)
    while True:
        line = f.readline()
        if not line:
            time.sleep(interval)
            continue
        yield line


def is_comment(line):
    return line[0] == '#' or line.strip()[0:2] == '//'


def make_keys(data, key_names):

    all_keys = set()
    for item in data:
        stack = [([], item)]
        while len(stack) > 0:
            path, item = stack.pop()
            for key in item:
                if type(item[key]) == dict:
                    stack.append((path + [key], item[key]))
                else:
                    all_keys.add('.'.join(path + [key]))

    keys = []
    for name in key_names:
        if '*' in name:
            pattern = re.compile(re.sub('*', '.*', name))
            for k in all_keys:
                if pattern.fullmatch(k):
                    keys.append(k)
        elif name in all_keys:
            keys.append(name)
        else:
            click.echo('[Warn] unknown key {}, was ignored'.format(name), err=True)

    return keys


def get(item, key, index=0):
    if key == 'index':
        return index
    if key in item:
        return item[key]
    if '.' in key:
        subs = key.split('.')
        value = item
        for sub in subs:
            if sub not in value:
                return None
            value = value[sub]
        return value
    return None


@click.command()
@click.option('--port', '-P', default=8097, type=int, help='the server')
@click.option('-x', default='index', help='x-axis key; `index` is a special key')
@click.option('-y', required=True,
              help='y-axis keys; multiple keys can be specified by comma-split')
@click.option('--title')
@click.option('--x-label')
@click.option('--y-label')
@click.option('--y-min', default=None, type=float)
@click.option('--y-max', default=None, type=float)
@click.option('--smoothing', default=0, type=int, help='Simple Moving Average with recent N items')
@click.option('--interval', '-i', default=-1, type=int,
              help='plot auto-refresh interval (sec); when negative, disabled (default=-1)')
@click.argument('logpath')
def main(port, x, y, title,
         x_label, y_label,
         y_min, y_max,
         smoothing, interval, logpath):

    data = [json.loads(line) for line in open(logpath) if not is_comment(line)]
    title = title if title else logpath

    x_key = x
    y_keys = make_keys(data, y.split(','))

    def align(xs, ys):
        assert(len(xs) == len(ys))
        xs_ = []
        ys_ = []
        for i in range(len(xs)):
            if None in xs[i]:
                continue
            if None in ys[i]:
                continue
            xs_.append(xs[i])
            ys_.append(ys[i])
        return xs_, ys_

    X = [[get(item, x_key, i)] * len(y_keys) for i, item in enumerate(data)]
    Y = [[get(item, key, i) for key in y_keys] for i, item in enumerate(data)]
    X, Y = align(X, Y)
    X = numpy.array(X)
    Y = numpy.array(Y)

    if smoothing > 0:
        last_Y = Y[-smoothing:, :]
        smoothed_Y = Y + 0.0
        for i in range(min(smoothing, len(Y))):
            smoothed_Y[i] = numpy.mean(Y[:i + 1, :], axis=0)
        for i in range(smoothing, len(Y)):
            smoothed_Y[i] = numpy.mean(Y[i - smoothing + 1: i + 1, :], axis=0)
        Y = smoothed_Y

    plot = Visdom(port=port)
    opts = {'legend': y_keys}
    if title:
        opts['title'] = title
    if x_label:
        opts['xlabel'] = x_label
    if y_label:
        opts['ylabel'] = y_label
    if y_max:
        opts['ytickmax'] = y_max

    # range
    if y_min or y_max:
        opts['ytickmin'] = y_min or Y.min()
        opts['ytickmax'] = y_max or Y.max()

    win = plot.line(Y, X, opts=opts)

    if interval <= 0:
        return

    N = len(data)
    with open(logpath, 'r') as f:
        for line in follow(f, interval):
            N += 1
            data = json.loads(line)
            X = [[get(data, x_key, N) for key in y_keys]]
            Y = [[get(data, key, N) for key in y_keys]]
            X, Y = align(X, Y)
            X = numpy.array(X)
            Y = numpy.array(Y)
            if smoothing > 0:
                last_Y = numpy.vstack((last_Y, Y))[-smoothing:, :]
                Y = numpy.mean(last_Y, axis=0).reshape((1, -1))
            plot.line(Y, X, win=win, update='append')
            print('append', data)


if __name__ == '__main__':
    main()
