"""页面解析能力评测的指标计算模块。

指标说明：
- 坐标匹配：预测框与真实框的 IoU > 0.5 视为匹配成功
- 精确率 (Precision) = TP / (TP + FP)
- 召回率 (Recall) = TP / (TP + FN)
- 描述相似度：对坐标匹配的元素对，使用 SequenceMatcher 计算 description 的相似度
"""

from difflib import SequenceMatcher
from typing import List, Dict, Tuple, Optional


def compute_iou(box1: List[int], box2: List[int]) -> float:
    """计算两个边界框 [x1, y1, x2, y2] 的 IoU。"""
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])

    inter_w = max(0, x2 - x1)
    inter_h = max(0, y2 - y1)
    inter_area = inter_w * inter_h

    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union_area = area1 + area2 - inter_area

    if union_area <= 0:
        return 0.0
    return inter_area / union_area


def compute_description_similarity(desc1: str, desc2: str) -> float:
    """使用 SequenceMatcher 计算两个描述文本的相似度。"""
    return SequenceMatcher(None, desc1.lower(), desc2.lower()).ratio()


def match_elements(
    pred_elements: List[Dict],
    gt_elements: List[Dict],
    iou_threshold: float = 0.5,
) -> Tuple[List[Tuple[int, int, float]], List[int], List[int]]:
    """贪心匹配预测元素与真实元素。

    返回:
        matched_pairs: [(pred_idx, gt_idx, iou), ...] 匹配成功的元素对
        unmatched_pred: 未匹配的预测元素索引列表
        unmatched_gt: 未匹配的真实元素索引列表
    """
    # 计算所有预测-真实对的 IoU
    pairs = []
    for i, pred in enumerate(pred_elements):
        for j, gt in enumerate(gt_elements):
            iou = compute_iou(pred["coordinates"], gt["coordinates"])
            if iou >= iou_threshold:
                pairs.append((iou, i, j))

    # 按 IoU 降序贪心匹配
    pairs.sort(key=lambda x: x[0], reverse=True)

    matched_pred = set()
    matched_gt = set()
    matched_pairs = []

    for iou, i, j in pairs:
        if i in matched_pred or j in matched_gt:
            continue
        matched_pred.add(i)
        matched_gt.add(j)
        matched_pairs.append((i, j, iou))

    unmatched_pred = [i for i in range(len(pred_elements)) if i not in matched_pred]
    unmatched_gt = [j for j in range(len(gt_elements)) if j not in matched_gt]

    return matched_pairs, unmatched_pred, unmatched_gt


def evaluate_sample(
    pred_elements: List[Dict],
    gt_elements: List[Dict],
    iou_threshold: float = 0.5,
) -> Dict:
    """评测单条样本，返回精确率、召回率和描述相似度。

    返回:
        {
            "precision": float,
            "recall": float,
            "tp": int,
            "fp": int,
            "fn": int,
            "description_similarities": [float, ...],
            "avg_description_similarity": float,
        }
    """
    matched_pairs, unmatched_pred, unmatched_gt = match_elements(
        pred_elements, gt_elements, iou_threshold
    )

    tp = len(matched_pairs)
    fp = len(unmatched_pred)
    fn = len(unmatched_gt)

    precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0

    # 计算匹配对的描述相似度
    desc_similarities = []
    for pred_idx, gt_idx, _ in matched_pairs:
        pred_desc = pred_elements[pred_idx].get("description", "")
        gt_desc = gt_elements[gt_idx].get("description", "")
        sim = compute_description_similarity(pred_desc, gt_desc)
        desc_similarities.append(sim)

    avg_desc_sim = sum(desc_similarities) / len(desc_similarities) if desc_similarities else 0.0

    return {
        "precision": precision,
        "recall": recall,
        "tp": tp,
        "fp": fp,
        "fn": fn,
        "description_similarities": desc_similarities,
        "avg_description_similarity": avg_desc_sim,
    }


def evaluate_dataset(
    predictions: List[Dict],
    ground_truths: List[Dict],
    iou_threshold: float = 0.5,
) -> Dict:
    """评测整个数据集，返回聚合指标。

    predictions 和 ground_truths 中的每个条目应包含 "elements" 字段，
    两者通过 "image_path" 对齐。

    返回:
        {
            "num_samples": int,
            "overall_precision": float,
            "overall_recall": float,
            "overall_f1": float,
            "avg_description_similarity": float,
            "per_sample": [Dict, ...],
        }
    """
    gt_map = {item["image_path"]: item for item in ground_truths}

    total_tp = 0
    total_fp = 0
    total_fn = 0
    all_desc_similarities = []
    per_sample_results = []
    matched_count = 0

    for pred in predictions:
        image_path = pred["image_path"]
        gt = gt_map.get(image_path)

        if gt is None:
            print(f"Warning: 未找到 {image_path} 对应的 ground truth，跳过")
            continue

        pred_elements = pred.get("elements", [])
        gt_elements = gt.get("elements", [])

        result = evaluate_sample(pred_elements, gt_elements, iou_threshold)
        result["image_path"] = image_path
        per_sample_results.append(result)

        total_tp += result["tp"]
        total_fp += result["fp"]
        total_fn += result["fn"]
        all_desc_similarities.extend(result["description_similarities"])
        matched_count += len(result["description_similarities"])

    overall_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0
    overall_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
    overall_f1 = (
        2 * overall_precision * overall_recall / (overall_precision + overall_recall)
        if (overall_precision + overall_recall) > 0
        else 0.0
    )
    avg_desc_sim = sum(all_desc_similarities) / len(all_desc_similarities) if all_desc_similarities else 0.0

    return {
        "num_samples": len(per_sample_results),
        "total_tp": total_tp,
        "total_fp": total_fp,
        "total_fn": total_fn,
        "overall_precision": overall_precision,
        "overall_recall": overall_recall,
        "overall_f1": overall_f1,
        "avg_description_similarity": avg_desc_sim,
        "per_sample": per_sample_results,
    }