import os, re, json, math
import xml.etree.ElementTree as ET
from collections import defaultdict
from glob import glob
import sys
sys.setrecursionlimit(50000)

# ---------- 0. 公共工具 ----------

def verify_ui_elements_exist_adb(verified_dict, last_xml_path, click_point=None, penultimate_xml_path=None):
    def parse_xml(xml_path):
        try:
            tree = ET.parse(xml_path)
        except Exception:
            return []
        root = tree.getroot()

        all_pairs = []
        def walk(p):
            for n in p:
                all_pairs.append((p, n))
                walk(n)
        walk(root)
        return all_pairs

    def match_element_in_xml(groups, all_pairs):
        # print(groups)
        for group in groups:  # or-group
            group_ok = True
            for pair in group:  # and-group
                t_self = pair.get("self", {})
                t_parent = pair.get("parent", {})

                # 如果 parent 是空，就不做 parent 比较
                parent_match_required = bool(t_parent)

                pair_ok = False
                for p_node, c_node in all_pairs:
                    self_match = all(t_self.get(k) == c_node.attrib.get(k, "") for k in t_self)
                    parent_match = all(t_parent.get(k) == p_node.attrib.get(k, "") for k in t_parent) if parent_match_required else True

                    if self_match and parent_match:
                        pair_ok = True
                        break

                if not pair_ok:
                    group_ok = False
                    break

            if group_ok:
                return True
        return False

    def match_click_on_element(groups, all_pairs, click_point):
        if not click_point:
            return False

        cx, cy = click_point
        # print(click_point)
        for group in groups:
            group_ok = True
            for pair in group:
                t_self = pair.get("self", {})
                matched = False
                for p_node, node in all_pairs: 
                    if all(t_self.get(k) == node.attrib.get(k, "") for k in t_self):
                        # 提取边界
                        try:
                            bounds = p_node.attrib.get("bounds")  # 格式: [x1,y1][x2,y2]   # 这里使用父元素坐标框
                            coords = list(map(int, re.findall(r"\d+", bounds)))
                            if len(coords) == 4:
                                x1, y1, x2, y2 = coords
                                if x1 <= cx <= x2 and y1 <= cy <= y2:
                                    matched = True
                                    break
                        except:
                            pass
                if not matched:
                    group_ok = False
                    break

            if group_ok:
                return True
        return False

    # === 主逻辑 ===
    result = True

    if "last_page" in verified_dict:
        all_pairs = parse_xml(last_xml_path)
        if not match_element_in_xml(verified_dict["last_page"], all_pairs):
            result = False

    if "penultimate_action" in verified_dict:
        all_pairs = parse_xml(penultimate_xml_path)
        if not match_click_on_element(verified_dict["penultimate_action"], all_pairs, click_point):
            result = False

    return result




def load_element_pairs(element_path):
    pairs = {}
    with open(element_path, "r", encoding="utf-8") as f:
        for line in f:
            row = json.loads(line)
            pairs[row["instruction"]] = row["verified_elements"]
    return pairs


def analyze_run(run_path, element_pairs):
    """
    返回一次运行的各项指标（字典），并负责把 model_success 重算进样本
    """
    # ---------- 1. 重新判定 model_success ----------
    data = []
    with open(os.path.join(run_path, "result.jsonl"), "r", encoding="utf-8") as f:
        for line in f:
            data.append(json.loads(line))

    for d in data:
        instr = d["instruction"]
        trace_folder = os.path.join(run_path, "trace", instr)
        xml_files = glob(os.path.join(trace_folder, "*.xml"))
        if not xml_files:
            continue

        sorted_xmls = sorted(xml_files, key=lambda x: int(re.findall(r"(\d+)", os.path.basename(x))[0]))
        last_xml = sorted_xmls[-1]
        penultimate_xml = sorted_xmls[-2] if len(sorted_xmls) >= 2 else None

        ele = element_pairs.get(instr, [])
        if 'penultimate_action' in str(ele):
            json_path = os.path.join(trace_folder, "data.json")
            with open(json_path, "r", encoding="utf-8") as f:
                trace_log = json.load(f)
            if len(trace_log['steps']) < 2:
                verified = 0
                continue
            else:
                click_point = trace_log['steps'][-2]['pred_action']['action_param']
                if not (type(click_point) == list and len(click_point) == 2):
                    verified = 0
                    continue
            resolution = trace_log['screen_resolution']

            if any(x in run_path for x in ['uitars', 'agentcpm','atlas', 'gpt-4']):
                click_point = [click_point[0] /1000 * resolution[0], click_point[1] /1000 * resolution[1]]
            elif any(x in run_path for x in ['mimo', 'qwen']):
                click_point = [click_point[0], click_point[1]]
            elif any(x in run_path for x in ['tongui', 'aguvis']):
                click_point = [click_point[0] * resolution[0], click_point[1] * resolution[1]]
            else:
                raise
            # click_point = [click_point[0] /1000 * resolution[0], click_point[1] /1000 * resolution[1]]

            verified = verify_ui_elements_exist_adb(ele, last_xml, click_point, penultimate_xml)
        else:
            verified = verify_ui_elements_exist_adb(ele, last_xml)


        d["model_success"] = verified

    # ---------- 2. 统计 ----------
    bucket   = defaultdict(lambda: {"succ": 0, "tot": 0})
    app_stat = defaultdict(lambda: {"succ": 0, "tot": 0})
    step_dur_sum = step_dur_cnt = 0
    ratio_sum = ratio_cnt = 0
    total_out_tok = 0
    fp_finish = fn_finish = 0
    aaaa = 0

    for d in data:
        succ   = bool(d["model_success"])
        gstep  = d.get("golden_steps", math.inf)
        mstep  = d.get("model_steps", 0)
        app    = d.get("app_name", "unknown")
        aaaa += mstep

        key = gstep if gstep <= 7 else "8+"
        bucket[key]["tot"]  += 1
        bucket[key]["succ"] += succ

        app_stat[app]["tot"]  += 1
        app_stat[app]["succ"] += succ

        total_out_tok += d.get("total_output_tokens", 0)

        if succ and gstep:
            ratio_sum += mstep / gstep
            ratio_cnt += 1

        for p, inf in zip(d.get("prepare_data_times", []), d.get("infer_times", [])):
            step_dur_sum += (p + inf)
            step_dur_cnt += 1

        if d.get("model_self_finish") and not succ:
            fp_finish += 1
        if not d.get("model_self_finish") and succ:
            fn_finish += 1

    step_dur_cnt = step_dur_cnt if step_dur_cnt else aaaa

    return {
        "bucket": bucket,
        "app_stat": app_stat,
        "total_success": sum(d["model_success"] for d in data),
        "N": len(data),
        "fp_finish": fp_finish,
        "fn_finish": fn_finish,
        "step_dur_sum": step_dur_sum,
        "step_dur_cnt": step_dur_cnt,
        "ratio_sum": ratio_sum,
        "ratio_cnt": ratio_cnt,
        "total_out_tok": total_out_tok / step_dur_cnt,
    }


def print_run(idx, stat):
    print(f"\n======== 第 {idx} 次运行 ========")
    # 1) 按 golden_steps 成功率
    for k in [3,4,5,6,7,"8+"]:
        s, t = stat["bucket"][k]["succ"], stat["bucket"][k]["tot"] or 1
        print(f"  golden_steps={k:<2}: {s}/{t} → {s/t:.2%}")
    tot_succ, N = stat["total_success"], stat["N"]
    print(f"  总成功率: {tot_succ}/{N} → {tot_succ/N:.2%}")

    # 2) finished 误报/漏报
    print(f"  误报 finished 概率: {stat['fp_finish']/N:.2%}")
    print(f"  漏报 finished 概率: {stat['fn_finish']/N:.2%}")

    # 3) 平均 step 耗时 & token
    avg_step_time = stat["step_dur_sum"] / stat["step_dur_cnt"] if stat["step_dur_cnt"] else 0
    avg_out_tok   = stat["total_out_tok"] / N
    avg_ratio     = stat["ratio_sum"] / stat["ratio_cnt"] if stat["ratio_cnt"] else 0
    print(f"  平均(prepare+infer)/step: {avg_step_time:.2f} s")
    print(f"  平均 completion_tokens : {avg_out_tok:.2f}")
    print(f"  成功样本步数比均值     : {avg_ratio:.3f}")

    # 4) 各 App 成功率
    print("  各 App 成功率:")
    for app, v in sorted(stat["app_stat"].items()):
        s, t = v["succ"], v["tot"]
        print(f"    {app:6s}: {s}/{t} → {s/t:.2%}")


# ---------- 3. 主入口 ----------
com = "your_path/online_eval_result/{}"
mapping = {
    'uitars7b_home': ['21451181-111111_uitars7b'],
}


paths = [com.format(p) for p in mapping["agentcpm_other"]]

element_path = "your_path/guiceval_online.jsonl"
element_pairs = load_element_pairs(element_path)


global_stat = defaultdict(float)   # 用来做加权汇总
app_union   = defaultdict(lambda: {"succ": 0, "tot": 0})
bucket_union= defaultdict(lambda: {"succ": 0, "tot": 0})

for i, p in enumerate(paths, 1):
    st = analyze_run(p, element_pairs)
    print_run(i, st)

    # ------- 汇总到全局 -------
    for k, v in st["bucket"].items():
        bucket_union[k]["succ"] += v["succ"]
        bucket_union[k]["tot"]  += v["tot"]

    for app, v in st["app_stat"].items():
        app_union[app]["succ"] += v["succ"]
        app_union[app]["tot"]  += v["tot"]

    for k in ["total_success","N","fp_finish","fn_finish",
              "step_dur_sum","step_dur_cnt","ratio_sum","ratio_cnt","total_out_tok"]:
        global_stat[k] += st[k]

# ---------- 4. 全局均值 ----------
print("\n======== 多次运行汇总（加权平均） ========")
for k in [3,4,5,6,7,"8+"]:
    s, t = bucket_union[k]["succ"], bucket_union[k]["tot"] or 1
    print(f"  golden_steps={k:<2}: {s}/{t} → {s/t:.2%}")
tot_succ, N = global_stat["total_success"], global_stat["N"]
print(f"  总成功率: {tot_succ}/{N} → {tot_succ/N:.2%}")

print(f"  误报 finished 概率: {global_stat['fp_finish']/N:.2%}")
print(f"  漏报 finished 概率: {global_stat['fn_finish']/N:.2%}")

avg_step_time = global_stat["step_dur_sum"] / global_stat["step_dur_cnt"] if global_stat["step_dur_cnt"] else 0
print(f"  平均(prepare+infer)/step: {avg_step_time:.2f} s")

avg_out_tok = global_stat["total_out_tok"] / len(paths)
print(f"  平均 completion_tokens : {avg_out_tok:.2f}")

avg_ratio = global_stat["ratio_sum"] / global_stat["ratio_cnt"] if global_stat["ratio_cnt"] else 0
print(f"  成功样本步数比均值     : {avg_ratio:.3f}")

print("  各 App 成功率:")
for app, v in sorted(app_union.items()):
    s, t = v["succ"], v["tot"]
    print(f"    {app:6s}: {s}/{t} → {s/t:.2%}")
