import json
import os
import numpy as np
import copy
import datetime
import sys
import re
from tqdm import tqdm
import ollama
import random
from collections import Counter

from torch.utils.data import Dataset

import nltk
from nltk.corpus.reader import VERB
import inflection

from datetime import datetime, timedelta, timezone

from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

_MANUAL_MATCHES = {
    "police": "police",
    "las": "las",
    "vegas": "vegas",
    "yes": "yes",
    "jeans": "jean",
    "hell's": "hell",
    "domino's": "domino",
    "morning": "morn",
    "clothes": "cloth",
    "are": "are",
    "riding": "ride",
    "leaves": "leaf",
    "dangerous": "danger",
    "clothing": "cloth",
    "texting": "text",
    "kiting": "kite",
    "firefighters": "firefight",
    "ties": "tie",
    "married": "married",
    "teething": "teeth",
    "gloves": "glove",
    "tennis": "tennis",
    "dining": "dine",
    "directions": "direct",
    "waves": "wave",
    "christmas": "christmas",
    "drives": "drive",
    "pudding": "pud",
    "coding": "code",
    "plating": "plate",
    "quantas": "quanta",
    "hornes": "horn",
    "graves": "grave",
    "mating": "mate",
    "paned": "pane",
    "alertness": "alert",
    "sunbathing": "sunbath",
    "tenning": "ten",
    "wetness": "wet",
    "urinating": "urine",
    "sickness": "sick",
    "braves": "brave",
    "firefighting": "firefight",
    "lenses": "lens",
    "reflections": "reflect",
    "backpackers": "backpack",
    "eatting": "eat",
    "designers": "design",
    "curiousity": "curious",
    "playfulness": "play",
    "blindness": "blind",
    "hawke": "hawk",
    "tomatoe": "tomato",
    "rodeoing": "rodeo",
    "brightness": "bright",
    "circuses": "circus",
    "skateboarders": "skateboard",
    "staring": "stare",
    "electronics": "electron",
    "electicity": "elect",
    "mountainous": "mountain",
    "socializing": "social",
    "hamburgers": "hamburg",
    "caves": "cave",
    "transitions": "transit",
    "wading": "wade",
    "creame": "cream",
    "toileting": "toilet",
    "sautee": "saute",
    "buildings": "build",
    "belongings": "belong",
    "stockings": "stock",
    "walle": "wall",
    "cumulis": "cumuli",
    "travelers": "travel",
    "conducter": "conduct",
    "browsing": "brows",
    "pooping": "poop",
    "haircutting": "haircut",
    "toppings": "top",
    "hearding": "heard",
    "sunblocker": "sunblock",
    "bases": "base",
    "markings": "mark",
    "mopeds": "mope",
    "kindergartener": "kindergarten",
    "pies": "pie",
    "scrapbooking": "scrapbook",
    "couponing": "coupon",
    "meetings": "meet",
    "elevators": "elev",
    "lowes": "low",
    "men's": "men",
    "childrens": "children",
    "shelves": "shelve",
    "paintings": "paint",
    "raines": "rain",
    "paring": "pare",
    "expressions": "express",
    "routes": "rout",
    "pease": "peas",
    "vastness": "vast",
    "awning": "awn",
    "boy's": "boy",
    "drunkenness": "drunken",
    "teasing": "teas",
    "conferences": "confer",
    "ripeness": "ripe",
    "suspenders": "suspend",
    "earnings": "earn",
    "reporters": "report",
    "kid's": "kid",
    "containers": "contain",
    "corgie": "corgi",
    "porche": "porch",
    "microwaves": "microwave",
    "batter's": "batter",
    "sadness": "sad",
    "apartments": "apart",
    "oxygenize": "oxygen",
    "striping": "stripe",
    "purring": "pure",
    "professionals": "profession",
    "piping": "pipe",
    "farmer's": "farmer",
    "potatoe": "potato",
    "emirates": "emir",
    "womens": "women",
    "veteran's": "veteran",
    "wilderness": "wilder",
    "propellers": "propel",
    "alpes": "alp",
    "charioteering": "chariot",
    "swining": "swine",
    "illness": "ill",
    "crepte": "crept",
    "adhesives": "adhesive",
    "regent's": "regent",
    "decorations": "decor",
    "rabbies": "rabbi",
    "overseas": "oversea",
    "travellers": "travel",
    "casings": "case",
    "smugness": "smug",
    "doves": "dove",
    "nationals": "nation",
    "mustange": "mustang",
    "ringe": "ring",
    "gondoliere": "gondolier",
    "vacationing": "vacate",
    "reminders": "remind",
    "baldness": "bald",
    "settings": "set",
    "glaced": "glace",
    "coniferous": "conifer",
    "revelations": "revel",
    "personals": "person",
    "daughter's": "daughter",
    "badness": "bad",
    "projections": "project",
    "polarizing": "polar",
    "vandalizers": "vandal",
    "minerals": "miner",
    "protesters": "protest",
    "controllers": "control",
    "weddings": "wed",
    "sometimes": "sometime",
    "earing": "ear",
}


class OKVQAStemmer:
    """Stemmer to match OKVQA v1.1 procedure."""

    def __init__(self):
        self._wordnet_lemmatizer = nltk.stem.WordNetLemmatizer()

    def stem(self, input_string):
        """Apply stemming."""
        word_and_pos = nltk.pos_tag(nltk.tokenize.word_tokenize(input_string))
        stemmed_words = []
        for w, p in word_and_pos:
            if w in _MANUAL_MATCHES:
                w = _MANUAL_MATCHES[w]
            elif w.endswith("ing"):
                w = self._wordnet_lemmatizer.lemmatize(w, VERB)
            elif p.startswith("NNS") or p.startswith("NNPS"):
                w = inflection.singularize(w)
            stemmed_words.append(w)
        return " ".join(stemmed_words)


stemmer = OKVQAStemmer()


def postprocess_ok_vqa_generation(predictions) -> str:
    prediction = re.split("Question|Answer|Short", predictions, 1)[0]
    prediction = re.split(", ", prediction, 1)[0]
    prediction_stem = stemmer.stem(prediction)
    return prediction_stem

class VQA:
    def __init__(self, annotation_file=None, question_file=None):
        """
        Constructor of VQA helper class for reading and visualizing questions and answers.
        :param annotation_file (str): location of VQA annotation file
        :return:
        """
        # load dataset
        self.dataset = {}
        self.questions = {}
        self.qa = {}
        self.qqa = {}
        self.imgToQA = {}
        if not annotation_file == None and not question_file == None:
            print("loading VQA annotations and questions into memory...")
            time_t = datetime.now(timezone.utc)
            dataset = json.load(open(annotation_file, "r"))
            questions = json.load(open(question_file, "r"))
            print(datetime.now(timezone.utc) - time_t)
            self.dataset = dataset
            self.questions = questions
            self.createIndex()

    def createIndex(self):
        # create index
        print("creating index...")
        imgToQA = {ann["image_id"]: [] for ann in self.dataset["annotations"]}
        qa = {ann["question_id"]: [] for ann in self.dataset["annotations"]}
        qqa = {ann["question_id"]: [] for ann in self.dataset["annotations"]}
        for ann in self.dataset["annotations"]:
            imgToQA[ann["image_id"]] += [ann]
            qa[ann["question_id"]] = ann
        for ques in self.questions["questions"]:
            qqa[ques["question_id"]] = ques
        print("index created!")

        # create class members
        self.qa = qa
        self.qqa = qqa
        self.imgToQA = imgToQA

    def info(self):
        """
        Print information about the VQA annotation file.
        :return:
        """
        for key, value in self.dataset["info"].items():
            print("%s: %s" % (key, value))

    def getQuesIds(self, imgIds=[], quesTypes=[], ansTypes=[]):
        """
        Get question ids that satisfy given filter conditions. default skips that filter
        :param 	imgIds    (int array)   : get question ids for given imgs
                        quesTypes (str array)   : get question ids for given question types
                        ansTypes  (str array)   : get question ids for given answer types
        :return:    ids   (int array)   : integer array of question ids
        """
        imgIds = imgIds if type(imgIds) == list else [imgIds]
        quesTypes = quesTypes if type(quesTypes) == list else [quesTypes]
        ansTypes = ansTypes if type(ansTypes) == list else [ansTypes]

        if len(imgIds) == len(quesTypes) == len(ansTypes) == 0:
            anns = self.dataset["annotations"]
        else:
            if not len(imgIds) == 0:
                anns = sum(
                    [self.imgToQA[imgId] for imgId in imgIds if imgId in self.imgToQA],
                    [],
                )
            else:
                anns = self.dataset["annotations"]
            anns = (
                anns
                if len(quesTypes) == 0
                else [ann for ann in anns if ann["question_type"] in quesTypes]
            )
            anns = (
                anns
                if len(ansTypes) == 0
                else [ann for ann in anns if ann["answer_type"] in ansTypes]
            )
        ids = [ann["question_id"] for ann in anns]
        return ids

    def getImgIds(self, quesIds=[], quesTypes=[], ansTypes=[]):
        """
         Get image ids that satisfy given filter conditions. default skips that filter
         :param quesIds   (int array)   : get image ids for given question ids
        quesTypes (str array)   : get image ids for given question types
        ansTypes  (str array)   : get image ids for given answer types
         :return: ids     (int array)   : integer array of image ids
        """
        quesIds = quesIds if type(quesIds) == list else [quesIds]
        quesTypes = quesTypes if type(quesTypes) == list else [quesTypes]
        ansTypes = ansTypes if type(ansTypes) == list else [ansTypes]

        if len(quesIds) == len(quesTypes) == len(ansTypes) == 0:
            anns = self.dataset["annotations"]
        else:
            if not len(quesIds) == 0:
                anns = sum(
                    [self.qa[quesId] for quesId in quesIds if quesId in self.qa], []
                )
            else:
                anns = self.dataset["annotations"]
            anns = (
                anns
                if len(quesTypes) == 0
                else [ann for ann in anns if ann["question_type"] in quesTypes]
            )
            anns = (
                anns
                if len(ansTypes) == 0
                else [ann for ann in anns if ann["answer_type"] in ansTypes]
            )
        ids = [ann["image_id"] for ann in anns]
        return ids

    def loadQA(self, ids=[]):
        """
        Load questions and answers with the specified question ids.
        :param ids (int array)       : integer ids specifying question ids
        :return: qa (object array)   : loaded qa objects
        """
        if type(ids) == list:
            return [self.qa[id] for id in ids]
        elif type(ids) == int:
            return [self.qa[ids]]

    def showQA(self, anns):
        """
        Display the specified annotations.
        :param anns (array of object): annotations to display
        :return: None
        """
        if len(anns) == 0:
            return 0
        for ann in anns:
            quesId = ann["question_id"]
            print("Question: %s" % (self.qqa[quesId]["question"]))
            for ans in ann["answers"]:
                print("Answer %d: %s" % (ans["answer_id"], ans["answer"]))

    def loadRes(self, resFile, quesFile):
        """
        Load result file and return a result object.
        :param   resFile (str)     : file name of result file
        :return: res (obj)         : result api object
        """
        res = VQA()
        res.questions = json.load(open(quesFile))
        res.dataset["info"] = copy.deepcopy(self.questions["info"])
        res.dataset["task_type"] = copy.deepcopy(self.questions["task_type"])
        res.dataset["data_type"] = copy.deepcopy(self.questions["data_type"])
        res.dataset["data_subtype"] = copy.deepcopy(self.questions["data_subtype"])
        res.dataset["license"] = copy.deepcopy(self.questions["license"])

        print("Loading and preparing results...     ")
        time_t = datetime.now(timezone.utc)
        anns = json.load(open(resFile))
        assert type(anns) == list, "results is not an array of objects"
        annsQuesIds = [ann["question_id"] for ann in anns]
        # print set of question ids that do not have corresponding annotations

        # assert set(annsQuesIds) == set(self.getQuesIds()), \
        # 'Results do not correspond to current VQA set. Either the results do not have predictions for all question ids in annotation file or there is atleast one question id that does not belong to the question ids in the annotation file.'
        for ann in anns:
            quesId = ann["question_id"]
            if res.dataset["task_type"] == "Multiple Choice":
                assert (
                    ann["answer"] in self.qqa[quesId]["multiple_choices"]
                ), "predicted answer is not one of the multiple choices"
            qaAnn = self.qa[quesId]
            ann["image_id"] = qaAnn["image_id"]
            ann["question_type"] = qaAnn["question_type"]
            if "answer_type" in ann:
                ann["answer_type"] = qaAnn["answer_type"]
        print(
            "DONE (t=%0.2fs)" % ((datetime.now(timezone.utc) - time_t).total_seconds())
        )

        res.dataset["annotations"] = anns
        res.createIndex()
        return res

class VQAEval:
    def __init__(self, vqa, vqaRes, n=2):
        self.n = n
        self.accuracy = {}
        self.evalQA = {}
        self.evalQuesType = {}
        self.evalAnsType = {}
        self.vqa = vqa
        self.vqaRes = vqaRes
        if not vqa is None and not vqaRes is None:
            self.params = {"question_id": vqaRes.getQuesIds()}
        self.contractions = {
            "aint": "ain't",
            "arent": "aren't",
            "cant": "can't",
            "couldve": "could've",
            "couldnt": "couldn't",
            "couldn'tve": "couldn't've",
            "couldnt've": "couldn't've",
            "didnt": "didn't",
            "doesnt": "doesn't",
            "dont": "don't",
            "hadnt": "hadn't",
            "hadnt've": "hadn't've",
            "hadn'tve": "hadn't've",
            "hasnt": "hasn't",
            "havent": "haven't",
            "hed": "he'd",
            "hed've": "he'd've",
            "he'dve": "he'd've",
            "hes": "he's",
            "howd": "how'd",
            "howll": "how'll",
            "hows": "how's",
            "Id've": "I'd've",
            "I'dve": "I'd've",
            "Im": "I'm",
            "Ive": "I've",
            "isnt": "isn't",
            "itd": "it'd",
            "itd've": "it'd've",
            "it'dve": "it'd've",
            "itll": "it'll",
            "let's": "let's",
            "maam": "ma'am",
            "mightnt": "mightn't",
            "mightnt've": "mightn't've",
            "mightn'tve": "mightn't've",
            "mightve": "might've",
            "mustnt": "mustn't",
            "mustve": "must've",
            "neednt": "needn't",
            "notve": "not've",
            "oclock": "o'clock",
            "oughtnt": "oughtn't",
            "ow's'at": "'ow's'at",
            "'ows'at": "'ow's'at",
            "'ow'sat": "'ow's'at",
            "shant": "shan't",
            "shed've": "she'd've",
            "she'dve": "she'd've",
            "she's": "she's",
            "shouldve": "should've",
            "shouldnt": "shouldn't",
            "shouldnt've": "shouldn't've",
            "shouldn'tve": "shouldn't've",
            "somebody'd": "somebodyd",
            "somebodyd've": "somebody'd've",
            "somebody'dve": "somebody'd've",
            "somebodyll": "somebody'll",
            "somebodys": "somebody's",
            "someoned": "someone'd",
            "someoned've": "someone'd've",
            "someone'dve": "someone'd've",
            "someonell": "someone'll",
            "someones": "someone's",
            "somethingd": "something'd",
            "somethingd've": "something'd've",
            "something'dve": "something'd've",
            "somethingll": "something'll",
            "thats": "that's",
            "thered": "there'd",
            "thered've": "there'd've",
            "there'dve": "there'd've",
            "therere": "there're",
            "theres": "there's",
            "theyd": "they'd",
            "theyd've": "they'd've",
            "they'dve": "they'd've",
            "theyll": "they'll",
            "theyre": "they're",
            "theyve": "they've",
            "twas": "'twas",
            "wasnt": "wasn't",
            "wed've": "we'd've",
            "we'dve": "we'd've",
            "weve": "we've",
            "werent": "weren't",
            "whatll": "what'll",
            "whatre": "what're",
            "whats": "what's",
            "whatve": "what've",
            "whens": "when's",
            "whered": "where'd",
            "wheres": "where's",
            "whereve": "where've",
            "whod": "who'd",
            "whod've": "who'd've",
            "who'dve": "who'd've",
            "wholl": "who'll",
            "whos": "who's",
            "whove": "who've",
            "whyll": "why'll",
            "whyre": "why're",
            "whys": "why's",
            "wont": "won't",
            "wouldve": "would've",
            "wouldnt": "wouldn't",
            "wouldnt've": "wouldn't've",
            "wouldn'tve": "wouldn't've",
            "yall": "y'all",
            "yall'll": "y'all'll",
            "y'allll": "y'all'll",
            "yall'd've": "y'all'd've",
            "y'alld've": "y'all'd've",
            "y'all'dve": "y'all'd've",
            "youd": "you'd",
            "youd've": "you'd've",
            "you'dve": "you'd've",
            "youll": "you'll",
            "youre": "you're",
            "youve": "you've",
        }
        self.manualMap = {
            "none": "0",
            "zero": "0",
            "one": "1",
            "two": "2",
            "three": "3",
            "four": "4",
            "five": "5",
            "six": "6",
            "seven": "7",
            "eight": "8",
            "nine": "9",
            "ten": "10",
        }
        self.articles = ["a", "an", "the", "be", "it\'s"]

        self.periodStrip = re.compile("(?!<=\\d)(\\.)(?!\\d)")
        self.commaStrip = re.compile("(\\d)(\\,)(\\d)")
        self.punct = [
            ";",
            r"/",
            "[",
            "]",
            '"',
            "{",
            "}",
            "(",
            ")",
            "=",
            "+",
            "\\",
            "_",
            "-",
            ">",
            "<",
            "@",
            "`",
            ",",
            "?",
            "!",
            ":",
        ]
        self.correct_questions = []

    def evaluate(self, quesIds=None):
        if quesIds == None:
            quesIds = [quesId for quesId in self.params["question_id"]]
        gts = {}
        res = {}
        for quesId in quesIds:
            gts[quesId] = self.vqa.qa[quesId]
            res[quesId] = self.vqaRes.qa[quesId]

        # =================================================
        # Compute accuracy
        # =================================================
        accQA = []
        accQuesType = {}
        accAnsType = {}
        print("computing accuracy")
        step = 0
        for quesId in quesIds:
            for ansDic in gts[quesId]["answers"]:
                ansDic["answer"] = ansDic["answer"].replace("\n", " ")
                ansDic["answer"] = ansDic["answer"].replace("\t", " ")
                ansDic["answer"] = ansDic["answer"].strip()
            resAns = res[quesId]["answer"]
            resAns = resAns.replace("\n", " ")
            resAns = resAns.replace("\t", " ")
            resAns = resAns.strip()
            resAns = self.processPunctuation(resAns)
            resAns = self.processDigitArticle(resAns)
            gtAcc = []

            for ansDic in gts[quesId]["answers"]:
                ansDic["answer"] = self.processPunctuation(ansDic["answer"])
                ansDic["answer"] = self.processDigitArticle(ansDic["answer"])

            for gtAnsDatum in gts[quesId]["answers"]:
                otherGTAns = [
                    item for item in gts[quesId]["answers"] if item != gtAnsDatum
                ]
                matchingAns = [item for item in otherGTAns if item["answer"] == resAns]
                acc = min(1, float(len(matchingAns)) / 3)
                gtAcc.append(acc)
            quesType = gts[quesId]["question_type"]
            ansType = (
                gts[quesId]["answer_type"] if "answer_type" in gts[quesId] else "other"
            )
            avgGTAcc = float(sum(gtAcc)) / len(gtAcc)
            accQA.append(avgGTAcc)
            if quesType not in accQuesType:
                accQuesType[quesType] = []
            accQuesType[quesType].append(avgGTAcc)
            if ansType not in accAnsType:
                accAnsType[ansType] = []
            accAnsType[ansType].append(avgGTAcc)
            self.setEvalQA(quesId, avgGTAcc)
            self.setEvalQuesType(quesId, quesType, avgGTAcc)
            self.setEvalAnsType(quesId, ansType, avgGTAcc)

            if step % 100 == 0:
                self.updateProgress(step / float(len(quesIds)))
            step = step + 1

            if avgGTAcc == 1:  # 如果准确率为 1，则认为该题目回答正确
                self.correct_questions.append(quesId)

        self.setAccuracy(accQA, accQuesType, accAnsType)
        print("Done computing accuracy")

        with open('/home/test/yxl/MCoT/okvqa/results/llama3.2-vision/correct_questions.txt', 'w') as f:
            for quesId in self.correct_questions:
                f.write(str(quesId) + '\n')

    def processPunctuation(self, inText):
        outText = inText
        for p in self.punct:
            if (p + " " in inText or " " + p in inText) or (
                re.search(self.commaStrip, inText) != None
            ):
                outText = outText.replace(p, "")
            else:
                outText = outText.replace(p, " ")
        outText = self.periodStrip.sub("", outText, re.UNICODE)
        return outText

    def processDigitArticle(self, inText):
        # 1. 初始化处理
        tempText = inText.lower().split()
        processed_words = []
        a_an_found = False  # 标记是否已找到第一个 'a' 或 'an'

        # 2. 遍历每个单词，只处理第一个 'a/an'
        for word in tempText:
            # 转换手动映射（如 "one" -> "1"）
            word = self.manualMap.setdefault(word, word)

            # 若已找到第一个 'a/an'，则收集后续单词直到遇到标点
            if a_an_found:
                # 检查是否是标点（利用已定义的 self.punct 集合）
                if (any(char in self.punct for char in word) or word in ['or', 'and']):
                    break  # 遇到标点则停止收集
                processed_words.append(word)
                continue

            # 检查是否是第一个 'a' 或 'an'
            if word in ['a', 'an', 'for'] and not a_an_found:
                a_an_found = True
                processed_words = []  # 清空之前内容，重新收集
                continue  # 跳过 'a/an' 本身，收集后续内容

            # 若未找到 'a/an'，按原逻辑过滤冠词
            if word not in self.articles:
                processed_words.append(word)

        # 3. 处理缩写词
        processed_text = " ".join(processed_words)
        processed_list = processed_text.split()
        for i, word in enumerate(processed_list):
            if word in self.contractions:
                processed_list[i] = self.contractions[word]
        processed_text = " ".join(processed_list)

        # 4. 若未找到 'a/an'，则按原逻辑过滤冠词
        if not a_an_found:
            outText = [word for word in tempText if word not in self.articles]
            processed_text = " ".join(outText)
            processed_text = " ".join([self.contractions.get(word, word) for word in processed_text.split()])

        return processed_text

    def setAccuracy(self, accQA, accQuesType, accAnsType):
        self.accuracy["overall"] = round(100 * float(sum(accQA)) / len(accQA), self.n)
        self.accuracy["perQuestionType"] = {
            quesType: round(
                100 * float(sum(accQuesType[quesType])) / len(accQuesType[quesType]),
                self.n,
            )
            for quesType in accQuesType
        }
        self.accuracy["perAnswerType"] = {
            ansType: round(
                100 * float(sum(accAnsType[ansType])) / len(accAnsType[ansType]), self.n
            )
            for ansType in accAnsType
        }

    def setEvalQA(self, quesId, acc):
        self.evalQA[quesId] = round(100 * acc, self.n)

    def setEvalQuesType(self, quesId, quesType, acc):
        if quesType not in self.evalQuesType:
            self.evalQuesType[quesType] = {}
        self.evalQuesType[quesType][quesId] = round(100 * acc, self.n)

    def setEvalAnsType(self, quesId, ansType, acc):
        if ansType not in self.evalAnsType:
            self.evalAnsType[ansType] = {}
        self.evalAnsType[ansType][quesId] = round(100 * acc, self.n)

    def updateProgress(self, progress):
        barLength = 20
        status = ""
        if isinstance(progress, int):
            progress = float(progress)
        if not isinstance(progress, float):
            progress = 0
            status = "error: progress var must be float\r\n"
        if progress < 0:
            progress = 0
            status = "Halt...\r\n"
        if progress >= 1:
            progress = 1
            status = "Done...\r\n"
        block = int(round(barLength * progress))
        text = "\rFinshed Percent: [{0}] {1}% {2}".format(
            "#" * block + "-" * (barLength - block), int(progress * 100), status
        )
        sys.stdout.write(text)
        sys.stdout.flush()


def extract_answer(output):
    pattern = r'The answer is[:\s]*(.*)'
    match = re.search(pattern, output)
    if match:
        answer = match.group(1).strip()
        period_pos = answer.find('.')
        if period_pos != -1:
            return answer[:period_pos].strip()
        return answer

    return output


class VQADataset(Dataset):
    def __init__(
        self, image_dir_path, question_path, annotations_path, is_train, dataset_name
    ):
        self.questions = json.load(open(question_path, "r"))["questions"]
        self.qid_to_idx = {q["question_id"]: i for i, q in enumerate(self.questions)}
        if annotations_path is not None:
            self.answers = json.load(open(annotations_path, "r"))["annotations"]
        else:
            self.answers = None

        self.image_dir_path = image_dir_path
        self.is_train = is_train
        self.dataset_name = dataset_name
        if self.dataset_name in {"vqav2", "ok_vqa"}:
            self.img_coco_split = self.image_dir_path.strip("/").split("/")[-1]
            assert self.img_coco_split in {"train2014", "val2014", "test2015"}

    def __len__(self):
        return len(self.questions)

    def get_img_path(self, question):
        if self.dataset_name in {"vqav2", "ok_vqa"}:
            return os.path.join(
                self.image_dir_path,
                f"COCO_{self.img_coco_split}_{question['image_id']:012d}.jpg"
                if self.is_train
                else f"COCO_{self.img_coco_split}_{question['image_id']:012d}.jpg",
            )
        elif self.dataset_name == "vizwiz":
            return os.path.join(self.image_dir_path, question["image_id"])
        elif self.dataset_name == "textvqa":
            return os.path.join(self.image_dir_path, f"{question['image_id']}.jpg")
        else:
            raise Exception(f"Unknown VQA dataset {self.dataset_name}")

    def __getitem__(self, idx):
        question = self.questions[idx]
        img_path = self.get_img_path(question)
        results = {
            "image_path": img_path,
            "question": question["question"],
            "question_id": question["question_id"],
        }
        if self.answers is not None:
            answers = self.answers[idx]
            results["answers"] = [a["answer"] for a in answers["answers"]]
        return results




def main():
    # 设置数据集路径
    train_image_dir_path = "/home/test/yxl/MCoT/data/okvqa/train2014"
    train_questions_json_path = "/home/test/yxl/MCoT/data/okvqa/OpenEnded_mscoco_train2014_questions.json"
    train_annotations_json_path = "/home/test/yxl/MCoT/data/okvqa/mscoco_train2014_annotations.json"
    test_image_dir_path = "/home/test/yxl/MCoT/data/okvqa/val2014"
    test_questions_json_path = "/home/test/yxl/MCoT/data/okvqa/OpenEnded_mscoco_val2014_questions.json"
    test_annotations_json_path = "/home/test/yxl/MCoT/data/okvqa/mscoco_val2014_annotations.json"

    # 创建数据集
    train_dataset = VQADataset(
        image_dir_path=train_image_dir_path,
        question_path=train_questions_json_path,
        annotations_path=train_annotations_json_path,
        is_train=True,
        dataset_name="ok_vqa",
    )

    test_dataset = VQADataset(
        image_dir_path=test_image_dir_path,
        question_path=test_questions_json_path,
        annotations_path=test_annotations_json_path,
        is_train=False,  # 测试集设为False
        dataset_name="ok_vqa",
    )

    random.seed(3)
    # 准备结果列表
    final_predictions = []

    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
    )

    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

    # 遍历数据集并调用模型
    for item in tqdm(test_dataset, desc="Processing dataset"):
        image_path = item["image_path"]
        question = item["question"]
        question_id = item["question_id"]

        prompt = f"please answer the question in the form of\' The answer is \'\n\nQuestion: {question}\n\n"


        # 构建消息
        messages = [{
            "role": "user",
            "content": prompt
        }]

        # 添加图像
        if os.path.exists(image_path):
            messages[0]["images"] = [image_path]

        vqa_eval = VQAEval(None, None)


        try:

            predictions = []
            all_outputs = []
            for _ in range(5):
                text = processor.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
                image_inputs, video_inputs = process_vision_info(messages)
                inputs = processor(
                    text=[text],
                    images=image_inputs,
                    videos=video_inputs,
                    padding=True,
                    return_tensors="pt",
                )
                inputs = inputs.to("cuda")

                generated_ids = model.generate(**inputs, max_new_tokens=128)
                generated_ids_trimmed = [
                    out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
                ]
                output_text = processor.batch_decode(
                    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
                )

                output = output_text[0]

                answer = extract_answer(output)
                answer = vqa_eval.processDigitArticle(answer)
                answer = vqa_eval.processPunctuation(answer)
                prediction = list(map(postprocess_ok_vqa_generation, [answer]))

                predictions.append(prediction[0])
                all_outputs.append(output)

            counter = Counter(predictions)
            final_prediction = counter.most_common(1)[0][0]



            # 保存结果
            final_predictions.append({
                "question_id": question_id,
                "question": question,
                "output": all_outputs,
                "answer": final_prediction
            })

        except Exception as e:
            print(f"Error processing question {question_id}: {e}")
            final_predictions.append({
                "question_id": question_id,
                "question": question,
                "answer": "ERROR"
            })

    # 保存结果到文件
    results_file = "/home/test/yxl/MCoT/okvqa/results/qwen2.5vl_test/SC.json"
    with open(results_file, "w") as f:
        json.dump(final_predictions, f, indent=4)



if __name__ == "__main__":
    main()
