import pandas as pd
import json
import glob

result = json.load(open("/home/qinyixin/workspace/TgMoE/Frozenbilm/results/T_MoENet_NEXT-QA.json"))
video_dir = "/mnt/hdd3/qinyixin/nextqa/video"

cols = pd.read_csv("/mnt/hdd3/qinyixin/FrozenBilm/NEXT-QA/val.csv").columns.to_list()
nextqa = pd.read_csv("/mnt/hdd3/qinyixin/FrozenBilm/NEXT-QA/val.csv").values
qid_to_vidid = {}
for it in nextqa:
    choices = [it[9 + idx] for idx in range(5)]
    answer = choices[it[6]]
    question = it[5]
    qid = it[7]
    vidid = str(it[1])
    vid_path = glob.glob(video_dir + "/*/"+ vidid + ".mp4")

    qid_to_vidid[str(qid)] = {"vid_path": vid_path, 
                              "choices": str(choices), 
                              "question": question,
                              "answer": answer
                              }

correct = []
for k, v in result.items():
    if v['acc']:
        correct.append(qid_to_vidid[k])

json.dump(correct, open("demo/T-MoENet_result.json", "w"))