import json
from pathlib import Path


def safe_set(value):
    if value is None:
        return set()
    if isinstance(value, list):
        return set(map(str, value))
    return {str(value)}


def prf1(tp: int, fp: int, fn: int):
    p = tp / (tp + fp) if tp + fp else 0.0
    r = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * p * r / (p + r) if p + r else 0.0
    return p, r, f1


def evaluate(dataset_path: str):
    data = json.loads(Path(dataset_path).read_text(encoding="utf-8"))
    # 数据格式示例：
    # [{"gold": {"laws": [...]}, "pred": {"laws": [...]}}]
    total_tp, total_fp, total_fn = 0, 0, 0
    for item in data:
        gold = safe_set(item["gold"].get("laws", []))
        pred = safe_set(item["pred"].get("laws", []))
        total_tp += len(gold & pred)
        total_fp += len(pred - gold)
        total_fn += len(gold - pred)
    p, r, f1 = prf1(total_tp, total_fp, total_fn)
    print(
        json.dumps(
            {"precision": round(p, 4), "recall": round(r, 4), "f1": round(f1, 4)},
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", required=True, help="标注评估集 json 路径")
    args = parser.parse_args()
    evaluate(args.dataset)