import argparse
import os
from log_operations import generators, generate_log


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary_dir", help="path to save summary logs", default="./summary")
    parser.add_argument("--step_num", type=int, help="step num", default=5)
    parser.add_argument("--tag_num", type=int, help="tag num", default=2)

    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = get_args()
    step_num = args.step_num
    tag_num = args.tag_num
    summary_dir = args.summary_dir
    if not os.path.exists(summary_dir):
        os.mkdir(summary_dir)
    
    # defined plugins, if you want to generate scalar only, you can define plugins = ['scalar']
    plugins = list(generators.keys())

    steps = [i for i in range(step_num)]
    tags = [f'tag_{i}' for i in range(tag_num)]
    generate_log(summary_dir, plugins, tags, steps)