data_handler.py

"""
数据加载与预处理模块
====================
负责加载 Parquet 文件，处理摘要列缺失，生成实验所需的数据记录。

如需更换数据源：
  - 修改 config.DATA_PATH 指向新文件
  - 确保新文件至少包含列: apa_reference_retracted, apa_reference_matched
  - 摘要列支持两种格式：
    (a) 单列 "abstract"（仅撤稿论文摘要）
    (b) 分列 "abstract_retracted" + "abstract_matched"（两篇都有）
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import pandas as pd

import config


# ──────────────────────────────────────────────
# 数据记录结构
# ──────────────────────────────────────────────

@dataclass
class PaperMetadata:
    """单篇论文的 OpenAlex 元数据。"""
    abstract: str | None = None
    doi: str | None = None
    pub_year: int | None = None
    citations: int | None = None
    topic: str | None = None


@dataclass
class PaperPairRecord:
    """一对匹配的论文记录（撤稿 + 对照）。"""
    row_index: int                        # 原始数据行号，便于追溯
    retracted_work_id: str | None         # 撤稿论文 ID（如果数据中有的话）
    control_work_id: str | None           # 对照论文 ID
    apa_retracted: str                    # 撤稿论文 APA 引用
    apa_matched: str                      # 对照论文 APA 引用
    abstract_retracted: str | None        # 撤稿论文摘要（可能为 None）
    abstract_matched: str | None          # 对照论文摘要（可能为 None）
    meta_retracted: PaperMetadata | None = None   # 撤稿论文元数据
    meta_matched: PaperMetadata | None = None     # 对照论文元数据


def load_data(path: str | None = None) -> list[PaperPairRecord]:
    """
    加载 Parquet 数据并返回 PaperPairRecord 列表。

    参数:
        path: Parquet 文件路径。默认使用 config.DATA_PATH。

    返回:
        PaperPairRecord 列表，每条代表一对撤稿-对照论文。
    """
    data_path = Path(path or config.DATA_PATH)
    if not data_path.exists():
        raise FileNotFoundError(
            f"数据文件不存在: {data_path.resolve()}\n"
            f"请将 retracted_control_matches_apa.parquet 放入 data/ 目录。"
        )

    df = pd.read_parquet(data_path, engine="pyarrow")

    # 验证必需列
    required_cols = {"apa_reference_retracted", "apa_reference_matched"}
    missing = required_cols - set(df.columns)
    if missing:
        raise ValueError(f"数据文件缺少必需列: {missing}")

    # 检测可选列
    has_retracted_id = "retracted_work_id" in df.columns
    has_control_id = "control_work_id" in df.columns

    # ---- 检测摘要列格式 ----
    # 格式 (b): 分列 abstract_retracted / abstract_matched（fetch_first_100.py 输出）
    # 格式 (a): 单列 abstract（旧格式，仅撤稿论文摘要）
    has_split_abstract = "abstract_retracted" in df.columns
    has_single_abstract = "abstract" in df.columns

    # 检测是否有 OpenAlex 元数据列（fetch_first_100.py 输出）
    META_FIELDS = ["doi", "pub_year", "citations", "topic"]
    has_metadata = "doi_retracted" in df.columns

    # 将 NaN 统一处理为 None（对所有可能含 NaN 的列）
    nullable_cols = ["abstract", "abstract_retracted", "abstract_matched"]
    if has_metadata:
        for f in META_FIELDS:
            nullable_cols += [f"{f}_retracted", f"{f}_matched"]
    for col in nullable_cols:
        if col in df.columns:
            df[col] = df[col].where(df[col].notna(), None)

    def _safe_val(val):
        """将 pandas NaN / NaT 统一为 None。"""
        if val is None:
            return None
        try:
            if pd.isna(val):
                return None
        except (TypeError, ValueError):
            pass
        return val

    def _build_meta(row, suffix: str) -> PaperMetadata | None:
        """从 DataFrame 行中提取元数据。"""
        if not has_metadata:
            return None
        doi = _safe_val(row.get(f"doi_{suffix}"))
        pub_year = _safe_val(row.get(f"pub_year_{suffix}"))
        citations = _safe_val(row.get(f"citations_{suffix}"))
        topic = _safe_val(row.get(f"topic_{suffix}"))
        # 如果全部为 None 则不创建对象
        if doi is None and pub_year is None and citations is None and topic is None:
            return None
        return PaperMetadata(
            doi=doi,
            pub_year=int(pub_year) if pub_year is not None else None,
            citations=int(citations) if citations is not None else None,
            topic=topic,
        )

    records: list[PaperPairRecord] = []
    for idx, row in df.iterrows():
        # 解析摘要：优先使用分列格式
        if has_split_abstract:
            abs_retracted = row.get("abstract_retracted")
            abs_matched = row.get("abstract_matched")
        elif has_single_abstract:
            abs_retracted = row.get("abstract")
            abs_matched = None
        else:
            abs_retracted = None
            abs_matched = None

        meta_ret = _build_meta(row, "retracted")
        meta_mat = _build_meta(row, "matched")

        # 把 abstract 也存进 metadata 对象（如果有的话）
        if abs_retracted and meta_ret:
            meta_ret.abstract = abs_retracted
        elif abs_retracted:
            meta_ret = PaperMetadata(abstract=abs_retracted)
        if abs_matched and meta_mat:
            meta_mat.abstract = abs_matched
        elif abs_matched:
            meta_mat = PaperMetadata(abstract=abs_matched)

        records.append(PaperPairRecord(
            row_index=int(idx),
            retracted_work_id=row.get("retracted_work_id") if has_retracted_id else None,
            control_work_id=row.get("control_work_id") if has_control_id else None,
            apa_retracted=str(row["apa_reference_retracted"]),
            apa_matched=str(row["apa_reference_matched"]),
            abstract_retracted=abs_retracted,
            abstract_matched=abs_matched,
            meta_retracted=meta_ret,
            meta_matched=meta_mat,
        ))

    print(f"[数据加载] 共加载 {len(records)} 对论文记录")
    if has_split_abstract or has_single_abstract:
        n_ret = sum(1 for r in records if r.abstract_retracted is not None)
        n_mat = sum(1 for r in records if r.abstract_matched is not None)
        print(f"[数据加载] 摘要: retracted={n_ret}, matched={n_mat}")
    else:
        print("[数据加载] 数据中无 abstract 列，Complex Prompt 将使用占位符")

    return records


def _format_metadata(meta: PaperMetadata | None) -> str:
    """将元数据格式化为可读文本片段。"""
    if meta is None:
        return ""
    parts = []
    if meta.doi:
        parts.append(f"DOI: {meta.doi}")
    if meta.pub_year is not None:
        parts.append(f"Publication year: {meta.pub_year}")
    if meta.citations is not None:
        parts.append(f"Cited by: {meta.citations}")
    if meta.topic:
        parts.append(f"Primary topic: {meta.topic}")
    return "\n".join(parts)


def build_prompt(
    template_name: str,
    apa: str,
    abstract: str | None = None,
    metadata: PaperMetadata | None = None,
) -> str:
    """
    根据模板名称和数据构建最终 Prompt。

    参数:
        template_name: "simple" 或 "complex"
        apa: APA 引用文本
        abstract: 摘要文本（仅 complex 模板需要）
        metadata: 论文元数据（仅 complex 模板使用）

    返回:
        格式化后的 Prompt 字符串
    """
    template = config.PROMPT_TEMPLATES[template_name]

    if template_name == "simple":
        return template.format(apa=apa)

    # Complex Prompt: 需要摘要 + 元数据
    if abstract is None or not isinstance(abstract, str) or abstract.strip() == "":
        abstract = "[Abstract not available]"

    meta_text = _format_metadata(metadata)
    if not meta_text:
        meta_text = "[Metadata not available]"

    return template.format(apa=apa, abstract=abstract, metadata=meta_text)


def get_prompt_types_for_data(records: list[PaperPairRecord]) -> list[str]:
    """
    根据数据是否含有摘要，决定可用的 Prompt 类型列表。

    - 如果有摘要: 返回 ["simple", "complex"]
    - 如果无摘要: 仅返回 ["simple"]（跳过 complex 避免无意义测试）
    """
    has_any_abstract = any(r.abstract_retracted is not None for r in records)
    if has_any_abstract:
        return ["simple", "complex"]
    return ["simple"]