import os
from datetime import datetime
from rich import print as pp


class DirectroyMaker:
    sub_dir_type = ['model', 'log', 'config']

    def __init__(self,
                 root,
                 save_model=True,
                 save_log=True,
                 save_config=True):
        self.root = os.path.expanduser(root)
        self.save_model = save_model
        self.save_log = save_log
        self.save_config = save_config
        self.detail_dir = None

    def create_dir(self, path):
        new_path = os.path.join(self.root,
                                self.detail_dir,
                                path)
        if not os.path.exists(new_path):
            os.makedirs(new_path)
        return new_path

    def experiments_dir_maker(self, args):
        if not os.path.exists(self.root):
            os.makedirs(self.root)
        now = datetime.now()
        time_idx = '_%s-%s-%s-%s-%s-%s' % (now.year,
                                           now.month,
                                           now.day,
                                           now.hour,
                                           now.minute,
                                           now.second)
        if args.experiments_name is None:
            self.detail_dir = f"{args.data_type}_{args.classifier_type}_{time_idx}"
        else:
            self.detail_dir = f"{args.experiments_name}_{args.tag}"
        if self.save_model:
            model_dir = self.create_dir(self.sub_dir_type[0])
            pp(f"[violet][Directory save_dir: {model_dir}][/violet]")
        if self.save_log:
            log_dir = self.create_dir(self.sub_dir_type[1])
            pp(f"[violet][Directory save_dir: {log_dir}][/violet]")
        if self.save_config:
            config_dir = self.create_dir(self.sub_dir_type[2])
            pp(f"[violet][Directory save_dir: {config_dir}][/violet]")

        return [model_dir, log_dir, config_dir]


