import argparse
from pathlib import Path
import pandas as pd
from rapidfuzz import fuzz
from contextlib import redirect_stdout, redirect_stderr
import time

fuzz_score_threshold = 80


def _normalize_cell(value: object) -> str:
    if pd.isna(value):
        return ""
    return str(value).strip()


def _load_rows(path: Path) -> list[tuple[str, str]]:
    df = pd.read_excel(path, dtype=str)
    if df.empty:
        return []

    first_two = df.iloc[:, :2].copy()
    # 避免 DataFrame.applymap 的 FutureWarning：按列使用 Series.map
    for col in first_two.columns:
        first_two[col] = first_two[col].map(_normalize_cell)
    first_two = first_two.replace("", pd.NA).dropna(how="all")
    rows: list[tuple[str, str]] = []
    seen: set[tuple[str, str]] = set()
    for item, text in first_two.itertuples(index=False, name=None):
        row = (
            "" if pd.isna(item) else str(item).strip(),
            "" if pd.isna(text) else str(text).strip(),
        )
        if row in seen:
            continue
        seen.add(row)
        rows.append(row)

    return rows


def _compare_impl(val_dir: Path, answer_dir: Path) -> None:
    val_dir = val_dir.resolve()
    answer_dir = answer_dir.resolve()

    overall_val = overall_answer = overall_matched = 0

    # 累计各“审查项”的全局统计
    overall_item_answer: dict[str, int] = {}
    overall_item_matched: dict[str, int] = {}
    overall_item_unmatched_answer: dict[str, int] = {}
    overall_item_unmatched_val: dict[str, int] = {}

    for val_file in sorted(val_dir.glob("*.xlsx")):
        answer_file = answer_dir / val_file.name
        if not answer_file.exists():
            print(f"Skip {val_file.name}: missing in answer")
            continue

        val_rows = _load_rows(val_file)
        answer_rows = _load_rows(answer_file)

        # Baseline: answer -> match val, consume val to keep 1-1, report leftover answers
        answer_counts: dict[str, int] = {}
        for item, _ in answer_rows:
            answer_counts[item] = answer_counts.get(item, 0) + 1

        val_buckets: dict[str, list[str]] = {}
        for item, text in val_rows:
            val_buckets.setdefault(item, []).append(text)

        matched_total = 0
        matched_by_item: dict[str, list[tuple[str, str, int]]] = {}
        unmatched_answer_by_item: dict[str, list[str]] = {}
        for item, ans_text in answer_rows:
            candidates = val_buckets.get(item, [])
            if not candidates:
                unmatched_answer_by_item.setdefault(item, []).append(ans_text)
                continue

            best_idx = -1
            best_score = -1
            for idx, cand in enumerate(candidates):
                ans_text = ans_text.strip()
                if cand is None or not isinstance(cand, str):
                    continue
                cand = cand.strip()
                score = max(
                    fuzz.partial_ratio(ans_text, cand),
                    fuzz.token_set_ratio(ans_text, cand),
                )
                if score > best_score:
                    best_score = score
                    best_idx = idx

            if best_score >= fuzz_score_threshold:
                matched_total += 1
                matched_val = candidates.pop(best_idx)
                matched_by_item.setdefault(item, []).append(
                    (ans_text, matched_val, best_score)
                )
            else:
                unmatched_answer_by_item.setdefault(item, []).append(ans_text)

        # remaining vals in buckets are unmatched
        unmatched_val_by_item: dict[str, list[str]] = {
            item: texts for item, texts in val_buckets.items() if texts
        }

        val_total = len(val_rows)
        answer_total = len(answer_rows)

        overall_val += val_total
        overall_answer += answer_total
        overall_matched += matched_total

        unmatched_val_count = sum(len(v) for v in unmatched_val_by_item.values())
        unmatched_answer_count = sum(len(v) for v in unmatched_answer_by_item.values())
        file_precision = (matched_total / val_total) if val_total != 0 else 0
        file_recall = (matched_total / answer_total) if answer_total != 0 else 0
        file_f1 = (
            (2 * file_precision * file_recall / (file_precision + file_recall))
            if (file_precision + file_recall)
            else 0
        )
        file_false_positive_rate = (
            (unmatched_val_count / val_total) if val_total != 0 else 0
        )

        # 累加到各“审查项”的全局统计
        for it, cnt in answer_counts.items():
            overall_item_answer[it] = overall_item_answer.get(it, 0) + cnt
        for it, lst in matched_by_item.items():
            overall_item_matched[it] = overall_item_matched.get(it, 0) + len(lst)
        for it, lst in unmatched_answer_by_item.items():
            overall_item_unmatched_answer[it] = overall_item_unmatched_answer.get(
                it, 0
            ) + len(lst)
        for it, lst in unmatched_val_by_item.items():
            overall_item_unmatched_val[it] = overall_item_unmatched_val.get(
                it, 0
            ) + len(lst)
        print("#" * 40)
        print(
            f"{val_file.name}: matched {matched_total} | val {val_total} | answer {answer_total} "
            f"| unmatched val {unmatched_val_count} | unmatched answer {unmatched_answer_count} | precision {file_precision:.2%} | recall {file_recall:.2%} | f1 {file_f1:.2%} | false_positive_rate {file_false_positive_rate:.2%}"
        )
        import json

        print(
            f"unmatched_val_by_item: {json.dumps(unmatched_val_by_item, ensure_ascii=False, indent=2)}"
        )
        for item in sorted(answer_counts):
            item_matches = matched_by_item.get(item, [])
            print(
                f"  审查项 {item}: matched {len(item_matches)} / {answer_counts[item]}"
            )
            # 匹配成功的结果
            # for ans_text, val_text, score in item_matches:
            # 	print(f"    {score}% | answer: {ans_text} | val: {val_text}")

            ua = unmatched_answer_by_item.get(item, [])
            if ua:
                print(f"    未匹配（answer 未被匹配）{len(ua)} 条：")
                for t in ua:
                    print(f"      answer: {t}")

            uv = unmatched_val_by_item.get(item, [])
            if uv:
                print(f"    未匹配（val 残留）{len(uv)} 条：")
                for t in uv:
                    print(f"      val: {t}")
        # break  # only first file for demo
    precision = overall_matched / overall_val if overall_val else 0
    recall = overall_matched / overall_answer if overall_answer else 0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0
    overall_false_positive_rate = (
        (overall_val - overall_matched) / overall_val if overall_val else 0
    )
    print(
        f"Overall: matched {overall_matched} | val {overall_val} | answer {overall_answer} | precision {precision:.2%} | recall {recall:.2%} | f1 {f1:.2%}"
    )

    # 按“审查项”的 overall 结果
    if overall_item_answer:
        print("#" * 40)
        print("Overall by item:")
        all_items = sorted(
            set(
                list(overall_item_answer.keys())
                + list(overall_item_matched.keys())
                + list(overall_item_unmatched_answer.keys())
                + list(overall_item_unmatched_val.keys())
            )
        )
        rows_by_item = []
        for it in all_items:
            ans = overall_item_answer.get(it, 0)
            mat = overall_item_matched.get(it, 0)
            u_ans = overall_item_unmatched_answer.get(it, 0)
            u_val = overall_item_unmatched_val.get(it, 0)
            item_precision = (mat / (mat + u_val)) if (mat + u_val) else 0
            acc = (mat / ans) if ans else 0
            item_f1 = (
                (2 * item_precision * acc / (item_precision + acc))
                if (item_precision + acc)
                else 0
            )
            item_false_positive_rate = u_val / (mat + u_val) if (mat + u_val) else 0
            rows_by_item.append(
                {
                    "审查项": it,
                    "大模型匹配上的不合格项": mat,
                    "合同所有不合格项": ans,
                    "大模型其他不合格项": u_val,
                    "大模型未匹配上的不合格项(C-B)": u_ans,
                    "查准率(B/B+D)": item_precision,
                    "查全率(B/C)": acc,
                    "F1": item_f1,
                    "误报率(D/B+D)": item_false_positive_rate,
                }
            )
            print(
                f"  审查项 {it}: matched {mat} / answer {ans} | unmatched val {u_val} | unmatched answer {u_ans} | precision {item_precision:.2%} | recall {acc:.2%} | f1 {item_f1:.2%}"
            )

        overall_by_item_df = pd.DataFrame(
            rows_by_item,
            columns=[
                "审查项",
                "大模型匹配上的不合格项",
                "合同所有不合格项",
                "大模型其他不合格项",
                "大模型未匹配上的不合格项(C-B)",
                "查准率(B/B+D)",
                "查全率(B/C)",
                "F1",
                "误报率(D/B+D)",
            ],
        )
        unmatched_val_total = sum(overall_item_unmatched_val.values())
        unmatched_answer_total = sum(overall_item_unmatched_answer.values())
        overall_precision = (
            overall_matched / (overall_matched + unmatched_val_total)
            if (overall_matched + unmatched_val_total)
            else 0
        )
        overall_f1 = (
            (2 * overall_precision * recall / (overall_precision + recall))
            if (overall_precision + recall)
            else 0
        )
        overall_invalid_rate = (
            unmatched_val_total / (overall_matched + unmatched_val_total)
            if (overall_matched + unmatched_val_total)
            else 0
        )
        overall_total_df = pd.DataFrame(
            [
                {
                    "审查项": "总体",
                    "大模型匹配上的不合格项": overall_matched,
                    "合同所有不合格项": overall_answer,
                    "大模型其他不合格项": unmatched_val_total,
                    "大模型未匹配上的不合格项(C-B)": unmatched_answer_total,
                    "查准率(B/B+D)": overall_precision,
                    "查全率(B/C)": recall,
                    "F1": overall_f1,
                    "误报率(D/B+D)": overall_invalid_rate,
                }
            ],
            columns=[
                "审查项",
                "大模型匹配上的不合格项",
                "合同所有不合格项",
                "大模型其他不合格项",
                "大模型未匹配上的不合格项(C-B)",
                "查准率(B/B+D)",
                "查全率(B/C)",
                "F1",
                "误报率(D/B+D)",
            ],
        )
        combined_df = pd.concat(
            [overall_by_item_df, overall_total_df], ignore_index=True
        )

        compare_dir_name = val_dir.name
        results_dir = Path(__file__).parent / "results"
        results_dir.mkdir(parents=True, exist_ok=True)
        output_excel = results_dir / f"合同审查结果_{compare_dir_name}.xlsx"
        with pd.ExcelWriter(output_excel, engine="openpyxl") as writer:
            combined_df.to_excel(writer, sheet_name="对比结果", index=False)
        print(
            f"Excel written to {output_excel}.\nEval Time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}"
        )


def compare(val_dir: Path, answer_dir: Path) -> None:
    _compare_impl(val_dir=val_dir, answer_dir=answer_dir)


def compare_with_log(
    val_dir: Path, answer_dir: Path, log_path: Path | None = None
) -> Path:
    val_dir = val_dir.resolve()
    if log_path is None:
        results_dir = Path(__file__).parent / "results"
        results_dir.mkdir(parents=True, exist_ok=True)
        log_path = results_dir / f"合同审查结果_{val_dir.name}.log"
    else:
        log_path = log_path.resolve()
        log_path.parent.mkdir(parents=True, exist_ok=True)

    with open(log_path, "w", encoding="utf-8") as f, redirect_stdout(
        f
    ), redirect_stderr(f):
        _compare_impl(val_dir=val_dir, answer_dir=answer_dir)

    return log_path


def _parse_args() -> argparse.Namespace:
    base = Path(__file__).parent
    parser = argparse.ArgumentParser(
        description="Compare extracted annotations with answers."
    )
    parser.add_argument(
        "--val-dir",
        type=Path,
        default=base / "jp-output-rj-base",
        help="Directory containing extracted val xlsx files.",
    )
    parser.add_argument(
        "--answer-dir",
        type=Path,
        default=base / "审查答案",
        help="Directory containing answer xlsx files.",
    )
    parser.add_argument(
        "--log-path",
        type=Path,
        default=None,
        help="Optional explicit log path. Defaults to results/合同审查结果_<val_dir_name>.log",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = _parse_args()
    final_log_path = compare_with_log(
        val_dir=args.val_dir,
        answer_dir=args.answer_dir,
        log_path=args.log_path,
    )
    print(f"Log written to {final_log_path}")
