VRCycling/Assets/Plotting/plot.py

import signal
import sys

from matplotlib import pyplot as plt
import asyncio
import os
import time

file_data = {}
axes = {}
labels = {}


def tryParseInt(value):
    if value is None or value == "":
        return value, False
    try:
        return int(value), True
    except ValueError:
        return value, False


def tryParseIntList(list):
    new_vars = []
    for i in list:
        y, success = tryParseInt(i)
        if not success:
            new_vars = None
            break
        new_vars.append(y)
    return new_vars


async def follow(path):
    with open(path, "r") as f:
        labels[path] = f.readline().split("\t")
        while True:
            line = f.readline()
            if not line:
                # print("nothing new")
                await asyncio.sleep(0.1)
                continue
            yield line
        print("returned follow")


def init_plt(paths):
    plt.ion()
    fig = plt.figure()

    nrows = 1
    ncolumns = 1
    index = 1

    while nrows * ncolumns < len(paths):
        nrows += 1
        if nrows * ncolumns >= len(paths):
            break
        ncolumns += 1

    for key in paths:
        axes[key] = fig.add_subplot(nrows, ncolumns, index)
        index += 1

    plt.subplots_adjust(hspace=0.5)

    return fig


def get_files():
    return ["plots/{}".format(p) for p in [f for f in os.listdir("plots") if f.endswith(".tsv")]]


async def read_file_data(path):
    print(f"reading data for {path}")
    lines = follow(path)
    xs = []
    ys = []

    async for line in lines:
        line_data = line.split("\t")

        x, success = tryParseInt(line_data[0])
        if success:
            xs.append(x / 1000)

            rest = tryParseIntList(line_data[1:])
            if rest is not None:
                if len(ys) == 0:
                    ys = [[i] for i in rest]
                else:
                    for i in range(len(rest)):
                        ys[i].append(rest[i])

                file_data[path] = (xs, ys)
                # print("added to file_data")

    print("read all lines")


async def plot_loop():
    print("before plot loop")
    while True:
        if len(axes) > 0:
            for key, ax in axes.items():
                if len(file_data) > 0:
                    xs, ys = file_data[key]
                    # print("reading file data")
                    ax.clear()
                    for i in range(len(ys)):
                        y = ys[i]
                        ax.plot(xs, y, label=labels[key][i + 1], linewidth=0.7)
                        ax.set_title(key)
                        ax.set_xlabel(labels[key][0])
                    left, right = ax.get_xlim()
                    if right - left > 10:
                        ax.set_xlim(right - 10, right)
            plt.legend()
            plt.pause(0.2)
            await asyncio.sleep(0.2)
            # await asyncio.sleep(1)


async def test():
    print("abc")


def signal_handler(signal, frame):
    loop.stop()
    sys.exit(0)


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    signal.signal(signal.SIGINT, signal_handler)

    paths = get_files()
    fig = init_plt(paths)

    asyncio.ensure_future(test())

    for p in paths:
        asyncio.ensure_future(read_file_data(p))

    asyncio.ensure_future(plot_loop())
    asyncio.ensure_future(test())

    loop.run_forever()