#!/usr/bin/env python3

import tqdm
import os
import time
import argparse
import pandas as pd
import tempfile
from pathlib import Path
from tensorboardX import SummaryWriter

parser = argparse.ArgumentParser()
parser.add_argument("--logdir", type=Path, default=".")
args, unk_args = parser.parse_known_args()

paths = sorted(args.logdir.rglob("*/**/*.csv"))
pdf = pd.DataFrame(paths, columns=["path"])
pdf["dir"] = pdf["path"].apply(lambda p: p.parent)

with tempfile.TemporaryDirectory() as tmpdir:
    pbar = tqdm.tqdm(pdf.groupby("dir"))
    for name, grp in pbar:
        pbar.set_description(f"Preprocessing {name}")
        writer = SummaryWriter(Path(tmpdir, name))
        for path in grp["path"]:
            df = pd.read_csv(path)
            df = df.sort_values("timestamp")
            rows = df.to_dict("records")
            for row in rows:
                step = None
                if "step" in row:
                    step = row["step"]
                timestamp = row["timestamp"]
                for key in set(row) - {"timestamp", "step"}:
                    try:
                        writer.add_scalar(key, row[key], step, timestamp)
                    except:
                        pass
    os.system(f"tensorboard --logdir {tmpdir} {' '.join(unk_args)}")
