import subprocess
subprocess.run(["pip", "install", "transformers==4.47.0"], check=True)
print("安装完成，继续执行后续代码...")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

import argparse, os, json, base64, re, httpx, torch, pytz
from openai import OpenAI
from collections import defaultdict
from datetime import datetime
from PIL import Image, ImageDraw
from tqdm import tqdm
import numpy as np
from typing import List

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

# ========== 工具函数 ==========
def encode_image(image_path: str) -> str:
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')


def write_jsonl(data: list, path: str):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'a', encoding='utf-8') as f:
        for d in data:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')


def prepare_data(system_prompt: str, json_path: str, imgs_path: str, vllm_format=False) -> list:
    items = []
    with open(json_path, 'r', encoding='utf-8') as f:
        for line in f:
            items.append(json.loads(line))

    label_data = []
    for data in items:
        img_file = os.path.join(imgs_path, data['img_path'])
        w, h = Image.open(img_file).size
        bbox = [
            int(data['bounds'][0] / w * 1000),
            int(data['bounds'][1] / h * 1000),
            int(data['bounds'][2] / w * 1000),
            int(data['bounds'][3] / h * 1000),
        ]
        if not vllm_format:
            # 默认格式
            query = system_prompt.format(data['instruction'])

            # vLLM 专用格式
        else:
            query = [{
                "role": "user",
                "content": [
                    {"type": 'image_url', "image_url": {"url": f"data:image/jpeg;base64,{encode_image(img_file)}"}},
                    {"type": 'text',      "text": system_prompt.format(data['instruction'])}
                ]
            }]

        label_data.append({
            'img_path': img_file,
            'query': query,
            'width': w,
            'height': h,
            'bbox': bbox,
            'device': data.get('device'),
            'data_type': data.get('data_type')
        })
    return label_data



def parse_output(text):
    try:
        data = json.loads(text)
        if all(k in data for k in ["x", "y", "width", "height"]):
            center_x = data["x"] + data["width"] // 2
            center_y = data["y"] + data["height"] // 2
            return [center_x, center_y]
    except json.JSONDecodeError:
        pass  # 继续执行正则解析

    pattern = r"(\d+\.\d+|\d+)[, ]+(\d+\.\d+|\d+)(?:[, ]+(\d+\.\d+|\d+)[, ]+(\d+\.\d+|\d+))?"
    matches = re.findall(pattern, text)

    results = []
    for match in matches:
        numbers = list(filter(None, match))  # 过滤掉空匹配项
        numbers = list(map(float, numbers))  # 转换为浮点数
        numbers = list(map(int, numbers))    # 取整（四舍五入改为 `round()`）

        if len(numbers) == 4:
            x1, y1, x2, y2 = numbers
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            results.append([center_x, center_y])  # 计算中心点
        elif len(numbers) == 2:
            results.append(numbers)
        else:
            results.append([0, 0])

    results = results[0] if len(results) != 0 else [0, 0]
    return results


def is_point_in_bbox(pt: list, bbox: list) -> bool:
    x, y = pt
    return bbox[0] <= x <= bbox[2] and bbox[1] <= y <= bbox[3]


def batch_data(data_list: list, bs: int):
    for i in range(0, len(data_list), bs):
        yield data_list[i:i+bs]

# ========== Transformers 模型实现 ==========
class TransformersVisionModel:
    def __init__(self, args):
        self.tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            args.model_path,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            device_map="auto",
            # quantization_config=BitsAndBytesConfig(load_in_8bit=True), # For INT8 quantization
            # quantization_config=BitsAndBytesConfig(load_in_4bit=True), # For INT4 quantization
        ).eval()


    def infer_batch(self, batch: list) -> list:
        assert len(batch) == 1

        query = batch[0]['query']
        image = Image.open(batch[0]['img_path']).convert("RGB")

        inputs = self.tokenizer.apply_chat_template(
            [{"role": "user", "image": image, "content": query}],
            add_generation_prompt=True,
            tokenize=True,
            return_tensors="pt",
            return_dict=True,
        ).to(self.model.device)
        # Generation parameters
        gen_kwargs = {
            "max_length": 4096,
            "do_sample": True,
            "top_k": 1,
        }
        with torch.no_grad():
            outputs = self.model.generate(**inputs, **gen_kwargs)
            outputs = outputs[:, inputs["input_ids"].shape[1]:]
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return [response]


# ========== vLLM API 实现 ==========
class VLLMVisionModel:
    def __init__(self, args):
        proxies = {"http://": args.proxy, "https://": args.proxy}
        client = httpx.Client(proxies=proxies)
        self.openai = OpenAI(
            base_url=args.vllm_url,
            api_key=args.api_key or "empty",
            http_client=client
        )
        self.deployment = args.deployment_name

    def infer_batch(self, batch: list) -> list:
        results = []
        for d in batch:
            try:
                res = self.openai.chat.completions.create(
                    model=self.deployment,
                    messages=d['query'],
                    frequency_penalty=1,
                    max_tokens=2048,
                    temperature=0.0,
                    top_p=0.7
                )
                results.append(res.choices[0].message.content)
            except Exception as e:
                print(f"⚠️ vLLM 调用失败: {e}")
                results.append(None)
        return results

# ========== 评估函数 ==========
def evaluate(model, args):
    system_prompt = "What's the proper procedure for \"{}\"? (with grounding)"
    data = prepare_data(
        system_prompt, args.json_path, args.img_path,
        vllm_format=bool(args.vllm_url)
    )
    all_scores, grouped = [], defaultdict(list)
    for batch in tqdm(batch_data(data, args.batch_size)):
        outputs = model.infer_batch(batch)
        for d, txt in zip(batch, outputs):
            d['response'] = txt
            d['predict'] = parse_output(txt)
            correct = is_point_in_bbox(d['predict'], d['bbox'])
            d['correct'] = correct
            # mask base64 in logs
            # if d['query'][0]['content'][0].get('image_url'):
            #     d['query'][0]['content'][0]['image_url'] = 'masked'
            d.pop('query')
            key = f"{d['device']}_{d['data_type']}"
            grouped[key].append(float(correct))
            all_scores.append(float(correct))
        write_jsonl(data=batch, path=args.result_path_file)
        for itm in batch:
            print(itm)
        
    args_dict = vars(args)
    write_jsonl([args_dict], args.result_path_file)
    stats = [{"item": k, "score": np.mean(v)} for k, v in grouped.items()]
    stats.append({"item": "all", "score": np.mean(all_scores)})
    write_jsonl(stats, args.result_path_file)
    for s in stats:
        print(s)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path',    '-m',    type=str, default='your_path/cogagent-9b-20241220')
    parser.add_argument('--json_path',     '-json', type=str, default='your_path/pub_version/grounding.jsonl')
    parser.add_argument('--img_path',      '-img',  type=str, default='your_path/GUI-CEval/img/grounding')
    parser.add_argument('--result_path',   '-o',    type=str, default='your_path/gui-ceval/grounding')
    parser.add_argument('--vllm_url',      '-url',  type=str, default=None)
    parser.add_argument('--deployment_name',        type=str, default='ui-tars-2b-sft')
    parser.add_argument('--api_key',       type=str, default=None)
    parser.add_argument('--proxy',         type=str, default='if needed')
    parser.add_argument('--batch_size',    type=int, default=1)
    args = parser.parse_args()

    # 构造输出文件路径
    now = datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y%m%d%H%M%S")
    fname = f"{now}_{os.path.basename(args.model_path)}_guiceval_grounding.jsonl"
    args.result_path_file = os.path.join(args.result_path, fname)

    # 选择模型实现
    vision_model = VLLMVisionModel(args) if args.vllm_url else TransformersVisionModel(args)

    evaluate(vision_model, args)