import os
import json


class Result_Logger:
    def __init__(self, path):
        self.log_dir = path
        self.log_path = os.path.join(path, "log.txt")
        if not os.path.exists(path):
            os.makedirs(path)
        print("Logging in:", path)

    def set_curr_dir(self, args):
        if args.few_samples > 0:
            trial_dir = "{}-samples-{}-seed".format(args.few_samples, args.seed)
        else:
            trial_dir = "{}-way-{}-shot-{}-seed".format(args.N, args.K, args.seed)
        trial_path = os.path.join(self.log_dir, trial_dir)
        self.curr_dir = trial_path
        if not os.path.exists(self.curr_dir):
            os.mkdir(self.curr_dir)

    def structure_save(self, structure, name):
        path = os.path.join(self.log_dir, name)
        with open(path, "w") as f:
            json.dump(structure, f, indent=4)
        print("dumping structure into", path)

    def _log_dict(self, info):
        s = ""
        for k, v in info.items():
            s += str(k) + " = " + str(v)
            s += ",\n"
        return s

    def log(self, info, seg=True):
        if isinstance(info, dict):
            s = self._log_dict(info)
        elif isinstance(info, str):
            s = info
        elif isinstance(info, list):
            s = str(info)
        with open(self.log_path, "a+") as f:
            f.writelines(s + "\n")
            if seg:
                f.writelines("=" * 20 + "\n")
        print(s)

    def log_structure(self, structure):
        assert isinstance(structure, dict)
        self.log(structure)

    def __call__(self, info, seg=True):
        self.log(info, seg)
