# read ./llava/backbone/results/kept_index.json
import json
# load tokenizer
from transformers import AutoTokenizer

if __name__ == "__main__":
    # parser
    import argparse
    import os
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", type=float, default=None)
    # cfg_ls
    parser.add_argument("--cfg_ls", type=float, nargs='+', default=None)
    # img_name
    parser.add_argument("--img_name", type=str, default="COCO_val2014_000000144305")
    # result path
    parser.add_argument("--result_path", type=str, default="./llava/backbone/results")
    args = parser.parse_args()

    model_name_or_path = "./checkpoints/llava-llama-2-7b-chat-lightning-preview"
    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path,
        padding_side="right",
        max_length=1024,
        use_fast=False,
    )# detokenize the kept index
    if args.cfg_ls is None and args.cfg is None:
        # list files from os.path.join(args.result_path, args.img_name, f'kept_index_{cfg}.json')
        # and get the cfg_ls
        import glob
        cfg_ls = []
        for file_path in glob.glob(os.path.join(args.result_path, args.img_name, f'kept_index_*.json')):
            cfg_ls.append(file_path.split('/')[-1].split('_')[-1][:-5])

    elif args.cfg is not None:
        cfg_ls = [args.cfg]
    else:
        cfg_ls = args.cfg_ls
    for cfg in cfg_ls:

        with open(os.path.join(args.result_path, args.img_name, f'kept_index_{cfg}.json')) as f:
            kept_index = [json.loads(line) for line in f]

        kept_index_detokenized = []
        for i in range(len(kept_index)):
            kept_index_detokenized.append([])
            for j in range(len(kept_index[i])):
                kept_index_detokenized[i].append(tokenizer.decode(kept_index[i][j]))
        # save to file
        with open(os.path.join(args.result_path, args.img_name, f"kept_index_detokenized_{cfg}.json"), 'w') as f:
            json.dump(kept_index_detokenized, f)
        print(f"saved to {os.path.join(args.result_path, args.img_name, f'kept_index_detokenized_{cfg}.json')}")



