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}...")

[((18, 22, 18, 46), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', ({(18, 36, 18, 45): 'self.path'}, {}), '(self.path)', False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((19, 21, 19, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import datetime, time\n'), ((32, 25, 32, 76), 'os.path.join', 'os.path.join', ({(32, 38, 32, 47): 'self.path', (32, 49, 32, 75): 'f"""{step:010d}_{score}.gif"""'}, {}), "(self.path, f'{step:010d}_{score}.gif')", False, 'import os\n'), ((33, 12, 33, 56), 'imageio.mimwrite', 'imageio.mimwrite', (), '', False, 'import imageio\n'), ((34, 12, 34, 32), 'pygifsicle.optimize', 'optimize', ({(34, 21, 34, 31): 'write_path'}, {}), '(write_path)', False, 'from pygifsicle import optimize\n'), ((12, 14, 12, 37), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime, time\n'), ((26, 33, 26, 44), 'time.time', 'time.time', ({}, {}), '()', False, 'import datetime, time\n')]