transform_jd_statement.py

from __future__ import annotations

import argparse
import csv
import re
from datetime import datetime
from pathlib import Path

import pandas as pd


MODEL_COLUMNS = [
    "ID",
    "Recency",
    "MntWines",
    "MntFruits",
    "MntMeatProducts",
    "MntFishProducts",
    "MntSweetProducts",
    "MntGoldProds",
    "NumDealsPurchases",
    "NumWebPurchases",
    "NumCatalogPurchases",
    "NumStorePurchases",
    "NumWebVisitsMonth",
    "AcceptedCmp1",
    "AcceptedCmp2",
    "AcceptedCmp3",
    "AcceptedCmp4",
    "AcceptedCmp5",
    "Complain",
]

CATEGORY_TO_MODEL_SPEND = {
    "酒类": "MntWines",
    "酒水饮料": "MntWines",
    "食品饮料": "MntFruits",
    "生鲜": "MntFruits",
    "蔬菜水果": "MntFruits",
    "水果": "MntFruits",
    "肉禽蛋": "MntMeatProducts",
    "肉类": "MntMeatProducts",
    "海鲜水产": "MntFishProducts",
    "水产": "MntFishProducts",
    "休闲食品": "MntSweetProducts",
    "糖果巧克力": "MntSweetProducts",
    "甜品": "MntSweetProducts",
    "珠宝首饰": "MntGoldProds",
    "黄金": "MntGoldProds",
}

TEXT_TO_MODEL_SPEND = [
    (re.compile(r"酒|白酒|啤酒|红酒|葡萄酒"), "MntWines"),
    (re.compile(r"水果|苹果|香蕉|橙|梨|桃|葡萄|草莓|蓝莓"), "MntFruits"),
    (re.compile(r"牛肉|猪肉|鸡肉|羊肉|肉|火腿|香肠"), "MntMeatProducts"),
    (re.compile(r"鱼|虾|蟹|海鲜|水产"), "MntFishProducts"),
    (re.compile(r"糖|巧克力|蛋糕|饼干|甜|零食"), "MntSweetProducts"),
    (re.compile(r"黄金|金条|金饰|首饰|珠宝"), "MntGoldProds"),
]

DISCOUNT_PATTERN = re.compile(r"优惠|券|红包|京豆|折扣|满减|促销")
REFUND_PATTERN = re.compile(r"退款|已全额退款|已退款")
STRONG_COMPLAIN_PATTERN = re.compile(r"退货|售后|投诉")


def find_table_start(lines: list[str]) -> int:
    for idx, line in enumerate(lines):
        if line.startswith("交易时间,商户名称,交易说明,金额"):
            return idx
    raise ValueError("Could not find the transaction table header.")


def parse_period_end(lines: list[str]) -> datetime | None:
    for line in lines[:15]:
        match = re.search(r"日期区间：\d{4}-\d{2}-\d{2}\s*至\s*(\d{4}-\d{2}-\d{2})", line)
        if match:
            return datetime.strptime(match.group(1), "%Y-%m-%d")
    return None


def parse_period_label(lines: list[str]) -> str | None:
    for line in lines[:15]:
        match = re.search(r"日期区间：(\d{4}-\d{2}-\d{2})\s*至\s*(\d{4}-\d{2}-\d{2})", line)
        if match:
            return f"{match.group(1)}_to_{match.group(2)}"
    return None


def parse_account_name(lines: list[str]) -> str | None:
    for line in lines[:15]:
        match = re.search(r"京东账号名：(.+)", line)
        if match:
            return match.group(1).strip()
    return None


def read_jd_statement(path: Path) -> tuple[pd.DataFrame, datetime | None]:
    lines = path.read_text(encoding="utf-8-sig").splitlines()
    table_start = find_table_start(lines)
    period_end = parse_period_end(lines)
    table_lines = [line.replace("\t", "") for line in lines[table_start:] if line.strip()]
    reader = csv.DictReader(table_lines)
    rows = list(reader)
    if not rows:
        raise ValueError("No transaction rows found.")
    return pd.DataFrame(rows), period_end


def map_transaction_to_model_col(category: str, description: str) -> str:
    category = str(category).strip()
    description = str(description)
    if category in CATEGORY_TO_MODEL_SPEND:
        return CATEGORY_TO_MODEL_SPEND[category]
    for pattern, target in TEXT_TO_MODEL_SPEND:
        if pattern.search(description):
            return target
    return "MntGoldProds"


def transform_jd_statement(input_path: Path, output_path: Path, customer_id: str | None) -> pd.DataFrame:
    lines = input_path.read_text(encoding="utf-8-sig").splitlines()
    raw, period_end = read_jd_statement(input_path)
    account_name = parse_account_name(lines) or input_path.stem
    period_label = parse_period_label(lines)
    resolved_customer_id = customer_id or (
        f"{account_name}_{period_label}" if period_label else account_name
    )
    raw["金额"] = pd.to_numeric(raw["金额"], errors="coerce").fillna(0)
    raw["交易时间"] = pd.to_datetime(raw["交易时间"], errors="coerce")

    successful_spending = raw[
        (raw["交易状态"] == "交易成功")
        & (raw["收/支"] == "支出")
        & (raw["金额"] > 0)
        & raw["交易时间"].notna()
    ].copy()
    if successful_spending.empty:
        raise ValueError("No successful spending transactions found.")

    as_of = period_end or successful_spending["交易时间"].max().to_pydatetime()
    latest_order = successful_spending["交易时间"].max().to_pydatetime()
    recency = max((as_of - latest_order).days, 0)

    aggregate = {col: 0 for col in MODEL_COLUMNS}
    aggregate["ID"] = resolved_customer_id
    aggregate["Recency"] = recency
    aggregate["NumWebPurchases"] = int(len(successful_spending))
    aggregate["NumCatalogPurchases"] = 0
    aggregate["NumStorePurchases"] = 0
    aggregate["NumWebVisitsMonth"] = int(max(aggregate["NumWebPurchases"], 1))
    raw_text = raw[["交易说明", "备注", "收/支", "交易状态"]].astype(str).agg(" ".join, axis=1)
    refund_count = int(raw_text.str.contains(REFUND_PATTERN).sum())
    strong_complain = bool(raw_text.str.contains(STRONG_COMPLAIN_PATTERN).any())
    refund_rate = refund_count / max(len(successful_spending), 1)
    aggregate["Complain"] = int(strong_complain or refund_rate >= 0.10)
    aggregate["NumDealsPurchases"] = int(
        successful_spending["交易说明"].astype(str).str.contains(DISCOUNT_PATTERN).sum()
    )

    for _, row in successful_spending.iterrows():
        target_col = map_transaction_to_model_col(row.get("交易分类", ""), row.get("交易说明", ""))
        aggregate[target_col] += float(row["金额"])

    # The statement does not expose campaign response history, so use neutral zeros.
    for col in ["AcceptedCmp1", "AcceptedCmp2", "AcceptedCmp3", "AcceptedCmp4", "AcceptedCmp5"]:
        aggregate[col] = 0

    result = pd.DataFrame([aggregate], columns=MODEL_COLUMNS)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.to_csv(output_path, index=False, encoding="utf-8-sig")

    audit = successful_spending[["交易时间", "交易分类", "交易说明", "金额"]].copy()
    audit["mapped_model_column"] = [
        map_transaction_to_model_col(category, description)
        for category, description in zip(audit["交易分类"], audit["交易说明"])
    ]
    audit.to_csv(output_path.with_name(output_path.stem + "_mapping_audit.csv"), index=False, encoding="utf-8-sig")
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description="Convert a JD transaction statement into model-ready features.")
    parser.add_argument("--input", required=True, help="Path to a JD transaction statement CSV.")
    parser.add_argument(
        "--output",
        default="outputs/jd_model_input.csv",
        help="Path for the model-ready output CSV.",
    )
    parser.add_argument(
        "--customer-id",
        default=None,
        help="Optional anonymized customer ID. Defaults to the JD account name in the export.",
    )
    args = parser.parse_args()

    result = transform_jd_statement(Path(args.input), Path(args.output), args.customer_id)
    print(f"Wrote model-ready JD features to {Path(args.output).resolve()}")
    print(result.to_string(index=False))


if __name__ == "__main__":
    main()