import numpy as np

from diffgro.utils.utils import print_b


class Logger:
    def __init__(self, title):
        self.title = title
        self.logs = {}

    def log(self, key, value, percent=False):
        if key not in self.logs:
            self.logs[key] = []
        if percent:
            value *= 100

        self.logs[key].append(value)

    def print(self):
        print_b("=" * 13 + f" {self.title} " + "=" * 13)
        for key, value in self.logs.items():
            avg = np.mean(value)
            std = np.std(value)
            print(f"\t{key}: {avg} +\- {std}")
        print_b("=" * 50)
