#!/usr/bin/env python

import glob
import os
import argparse


def get_sorted_chunk_ids(dirs):
    ids = []
    for d in dirs:
        for f in glob.glob(os.path.join(d, "game_*.gz")):
            if os.path.basename(f) == "game_000000.gz":
                ids.append((0, f))
            else:
                ids.append((int(
                    os.path.basename(f).split('.')[-2].split('_')[-1].lstrip(
                        "0")), f))
    ids.sort()
    return ids


def main(argv):
    a = get_sorted_chunk_ids([argv.input])
    for (i, f) in a:
        os.rename(
            f, os.path.join(argv.input,
                            "training.{}.gz".format(i + argv.base)))


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description=\
            'Rename files generated by lc0 to have training names.')
    argparser.add_argument('-i', '--input', type=str, help='input directory')
    argparser.add_argument('-b',
                           '--base',
                           type=int,
                           default=0,
                           help='base value for names')

    main(argparser.parse_args())
