import argparse

import numpy as np

from evalplus.data import get_human_eval_plus, get_mbpp_plus

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default="humaneval", choices=["mbpp", "humaneval"])
    parser.add_argument("--mini", action="store_true")
    args = parser.parse_args()

    print(f"Reporting stats for {args.dataset} dataset [ mini = {args.mini} ]")
    if args.dataset == "humaneval":
        data = get_human_eval_plus(mini=args.mini)
    elif args.dataset == "mbpp":
        data = get_mbpp_plus(mini=args.mini)

    sizes = np.array(
        [[len(inp["base_input"]), len(inp["plus_input"])] for inp in data.values()]
    )
    size_base = sizes[:, 0]
    print(f"{size_base.min() = }", f"{size_base.argmin() = }")
    print(f"{size_base.max() = }", f"{size_base.argmax() = }")
    print(f"{np.percentile(size_base, 50) = :.1f}")
    print(f"{size_base.mean() = :.1f}")

    size_plus = sizes[:, 1]
    size_plus += size_base
    print(f"{size_plus.min() = }", f"{size_plus.argmin() = }")
    print(f"{size_plus.max() = }", f"{size_plus.argmax() = }")
    print(f"{np.percentile(size_plus, 50) = :.1f}")
    print(f"{size_plus.mean() = :.1f}")