import os
import datetime, time

import imageio
from pygifsicle import optimize
from torch.utils.tensorboard import SummaryWriter


class LogManager:
    def __init__(self, env, id, experiment=None):
        self.id = id
        now = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")
        self.path = (
            f"./logs/{experiment}/{env}/{id}/{now}/"
            if experiment
            else f"./logs/{env}/{id}/{now}/"
        )
        self.writer = SummaryWriter(self.path)
        self.stamp = time.time()

    def write(self, scalar_dict, frames, step):
        for key, value in scalar_dict.items():
            self.writer.add_scalar(f"{self.id}/" + key, value, step)
            self.writer.add_scalar("all/" + key, value, step)
            if "score" in key:
                time_delta = int(time.time() - self.stamp)
                self.writer.add_scalar(f"{self.id}/{key}_per_time", value, time_delta)
                self.writer.add_scalar(f"all/{key}_per_time", value, time_delta)

        if len(frames) > 0:
            score = scalar_dict["score"]
            write_path = os.path.join(self.path, f"{step:010d}_{score}.gif")
            imageio.mimwrite(write_path, frames, fps=60)
            optimize(write_path)
            print(f"...Record episode to {write_path}...")

[((464, 488), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', (['self.path'], {}), '(self.path)\n', (477, 488), False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((510, 521), 'time.time', 'time.time', ([], {}), '()\n', (519, 521), False, 'import datetime, time\n'), ((1102, 1153), 'os.path.join', 'os.path.join', (['self.path', 'f"""{step:010d}_{score}.gif"""'], {}), "(self.path, f'{step:010d}_{score}.gif')\n", (1114, 1153), False, 'import os\n'), ((1166, 1210), 'imageio.mimwrite', 'imageio.mimwrite', (['write_path', 'frames'], {'fps': '(60)'}), '(write_path, frames, fps=60)\n', (1182, 1210), False, 'import imageio\n'), ((1223, 1243), 'pygifsicle.optimize', 'optimize', (['write_path'], {}), '(write_path)\n', (1231, 1243), False, 'from pygifsicle import optimize\n'), ((235, 258), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (256, 258), False, 'import datetime, time\n'), ((813, 824), 'time.time', 'time.time', ([], {}), '()\n', (822, 824), False, 'import datetime, time\n')]