#! /usr/bin/env python
import json

import benchml

log = benchml.log
benchml.readwrite.configure(use_ase=False)


def main(args):
    # Load datasets (as iterator)
    benchml.splits.synchronize(args.seed)
    data = benchml.data.BenchmarkData(
        root=args.data_folder, filter_fct=benchml.filters[args.filter]
    )
    # Compile models
    models = benchml.models.compile(args.groups)
    # Evaluate
    bench = benchml.benchmark.evaluate(data, models, log, verbose=args.verbose)
    json.dump(bench, open(args.output, "w"), indent=1, sort_keys=True)


if __name__ == "__main__":
    log.Connect()
    log.AddArg("data_folder", type=str, default="", help="Dataset folder")
    log.AddArg("groups", type=(list, str), default=[], help="Model groups")
    log.AddArg("filter", type=str, default="none", help="Dataset filter regex")
    log.AddArg("output", type=str, default="bench.json", help="Output json")
    log.AddArg("seed", type=int, default=0, help="Global random seed")
    log.AddArg("verbose", type="toggle", default=False, help="Toggle verbose output")
    log.AddArg(
        "list_transforms", type="toggle", default=False, help="List available transforms and quit"
    )
    log.AddArg(
        "list_collections", type="toggle", default=False, help="List model collections and quit"
    )
    args = log.Parse()
    if args.list_transforms:
        benchml.transforms.list_all(verbose=args.verbose)
        log.okquit()
    if args.list_collections:
        benchml.models.list_all(verbose=args.verbose)
        log.okquit()
    main(args)
