from data_generator import generate_data

import argparse


def parse_args():
    parser = argparse.ArgumentParser(description="Clear old dataset and generate new one")
    parser.add_argument("size", help="size of new dataset", type=int)
    parser.add_argument("depth", help="max_depth scope depth of generated equation, no less than 1", type=int)
    parser.add_argument("length", help="length of equation will be in range length/2..length", type=int)
    parser.add_argument("fraction", help="fraction of tex vocab to sample tokens from, float in range 0..1", type=float)
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    generate_data(examples_count=args.size, max_depth=args.depth, equation_length=args.length,
                  distribution_fraction=args.fraction)


if __name__ == "__main__":
    main()