from ast import Continue
import os, shutil, argparse

def reorder_and_copy(src_root, dst_root, keys):
    for dirpath, _, filenames in os.walk(src_root):

        # compute the a/b/c/d components from the directory path
        rel_dir = os.path.relpath(dirpath, src_root)           # e.g. "a/b/c/d"
        parts = rel_dir.split(os.sep)
        if len(parts) < 4 or len(parts) > 5:
            # skip directories that aren’t at least 4 deep
            Continue
        # find all files whose name contains the key substring
        matches = [f for f in filenames if any(key in f for key in keys)]
        if not matches:
            continue
        a, b, c, d = parts[:4]
        if "old" in a:
            continue
        print(os.path.join(a,b,c,d))

        # for each matching file, copy it into c/b/d/a/ keeping its original name
        for fname in matches:
            new_rel = os.path.join(c, b, d)
            dst_dir = os.path.join(dst_root, new_rel)
            os.makedirs(dst_dir, exist_ok=True)

            src_file = os.path.join(dirpath, fname)
            dst_file = os.path.join(dst_dir, fname)
            shutil.copy2(src_file, dst_file)
            # print(f"Copied {os.path.relpath(src_file, src_root)} → {dst_file}")

if __name__ == "__main__":
    keys = ["198.", "233.", "828", "1004", "1294"]
    reorder_and_copy("./playground/data/eval/mmbench", "./playground", keys)
    keys = ["828", "1004", "1294"]
    reorder_and_copy("./playground/quick_vis", "./playground/quick_vis_short", keys)
    print("Done.")
