import os
import argparse
from models import get_model_instance
import json
from tqdm import tqdm
from datetime import datetime

def rotate_abcd(options):
    """ 只轮转 A/B/C/D，其他选项保持不变 """
    abcd_keys = ['A', 'B', 'C', 'D']
    abcd_values = [options.get(k) for k in abcd_keys]

    # 轮转 A/B/C/D
    rotated_abcd = [abcd_values[-shift % 4:] + abcd_values[:-shift % 4] for shift in range(4)]

    # 更新 A/B/C/D 选项，其他选项不变
    for rotated in rotated_abcd:
        rotated_options = options.copy()
        for i, key in enumerate(abcd_keys):
            rotated_options[key] = rotated[i]
        yield rotated_options  # 每次 yield 一个轮换后的选项


def evaluate_from_jsonl(jsonl_path, topn, image_root, downscale_factor, model):
    correct, total = 0, 0
    abcd_keys = ['A', 'B', 'C', 'D']

    with open(jsonl_path, 'r', encoding='utf-8') as f:
        for line in tqdm(f):
            data = json.loads(line)
            if data['answer'] in [0, "0"]:
                total += 1
                continue

            question = data['question']
            options = {k: v for k, v in data.items() if k in abcd_keys and str(v).strip()}
            image_paths = [os.path.join(image_root, p) for p in data['img_path']] if isinstance(data['img_path'], list) else [os.path.join(image_root, data['img_path'])]
            if downscale_factor != 1:
                image_paths = [i_path.replace("/img/", f"/img_resized_{downscale_factor}/") for i_path in image_paths]
                
            passed_all = True
            for idx, rotated in enumerate(rotate_abcd(options)):
                prompt = f"{question}\n" + "\n".join([f"{k}: {v}" for k, v in rotated.items()])
                prompt += "\n请直接回复最合适的选项字母。"

                response = model.chat(prompt, image_paths)
                if not response:
                    passed_all = False
                    break

                pred = next((ch for ch in response.upper() if ch in rotated), "Invalid")
                correct_ans = abcd_keys[idx]
                if pred != correct_ans:
                    passed_all = False
                    break

            total += 1
            if passed_all:
                correct += 1
            if total == topn:
                break

    acc = correct / total if total else 0
    print(f"Accuracy: {acc:.2%}")
    return acc


def main(args):
    model = get_model_instance(args.model_name, args.ip)

    acc = evaluate_from_jsonl(
        jsonl_path=args.jsonl_path,
        topn=args.topn,
        image_root=args.image_root,
        downscale_factor=args.downscale_factor,
        model=model
    )

    result_path = args.result_path
    os.makedirs(result_path, exist_ok=True)
    result_path = os.path.join(result_path, args.model_name.split("/")[-1])
    log = f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] File: {args.jsonl_path} | Accuracy: {acc:.2%}\n"
    with open(f"{result_path}_eval_log.txt", "a", encoding="utf-8") as f:
        f.write(log)
    print(f"Logged result to {result_path}_eval_log.txt")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--jsonl_path', type=str, required=True)
    parser.add_argument('--topn', type=int, default=100000)
    parser.add_argument('--image_root', type=str, default="")
    parser.add_argument('--result_path', type=str, default='../results')
    parser.add_argument('--downscale_factor', type=int, default=1)
    parser.add_argument('--model_name', type=str, required=True)
    parser.add_argument('--gpt4o_choice', type=str, default="gpt-4o-2")
    parser.add_argument('--ip', type=str, default="")
    args = parser.parse_args()

    main(args)