"""生成大容量算术数据 —— 让模型没法靠"抄题目里的数字"蒙对。

为什么需要（第一版实测的问题）：
  只给 16 条算术样本时，模型学会了"输出题目里出现的数字"这个捷径：
      问 "1+1="  ->  答 "1+1=1"     （抄了 1）
      问 "2+2?"  ->  答 "2+2 is 2"  （抄了 2）
      问 "5+5?"  ->  答 "5+5 is 5"  （抄了 5）
  这是**数据量不足导致的伪规律**，不是加法没学会。

  修法：给**大量、操作数各不相同**的算术题。当题目里的数字千变万化而
  答案总等于两数之和时，"抄数字"这条捷径就失效了，模型只能学到真正的
  加法映射（至少在它能力范围内）。

  格式统一成 "<表达式> = <结果>."，答案唯一。

用法:
    python -m fahqgpt.make_math_data
"""
from __future__ import annotations

import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fahqgpt import config as C  # noqa: E402

TEMPLATES = ["{e}", "What is {e}?", "What is {e} =?", "{e} =", "Calculate {e}.", "{e}?"]


def main() -> int:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    random.seed(20260925)
    pairs: list[tuple[str, int]] = []

    # 加法：覆盖 0~100 的各种组合，尤其中小数字（模型最容易验证）
    for _ in range(1200):
        a = random.randint(0, 50)
        b = random.randint(0, 50)
        pairs.append((f"{a}+{b}", a + b))
    for _ in range(200):                     # 大一点的数
        a = random.randint(50, 200)
        b = random.randint(50, 200)
        pairs.append((f"{a}+{b}", a + b))
    # 交换律：b+a 也要有，避免模型只认 "小的在前"
    for _ in range(300):
        a = random.randint(0, 30)
        b = random.randint(0, 30)
        pairs.append((f"{b}+{a}", a + b))

    # 减法（结果非负）
    for _ in range(500):
        a = random.randint(1, 60)
        b = random.randint(0, a)
        pairs.append((f"{a}-{b}", a - b))

    # 乘法（小数字，乘法表范围）
    for _ in range(400):
        a = random.randint(1, 12)
        b = random.randint(1, 12)
        pairs.append((f"{a}*{b}", a * b))

    # 除法（整除，避免小数）
    for _ in range(200):
        b = random.randint(1, 12)
        q = random.randint(1, 12)
        pairs.append((f"{b*q}/{b}", q))

    # 去重（同表达式只留一个答案）
    seen: dict[str, int] = {}
    for e, r in pairs:
        seen[e] = r
    items = list(seen.items())
    random.shuffle(items)

    rows = []
    for e, r in items:
        # 答案用**纯数字**，不要 "12+1 = 13." 这种重复表达式的写法。
        # 原因：重复表达式等于给模型"抄表达式"的机会，它可以输出
        # "12+1 = 1" 这种半对半错的东西（第一版实测就是这样）。
        # 只给数字，它就无处可抄，必须真的算出结果。
        a = str(r)
        for t in TEMPLATES:
            rows.append({"user": t.format(e=e), "assistant": a})
    # 另外给一部分带句子的答案，避免模型只会输出裸数字
    for e, r in items[:400]:
        rows.append({"user": e, "assistant": f"The answer is {r}."})

    # 去重 (问题) -> 答案
    dedup = {}
    for r in rows:
        dedup[" ".join(r["user"].lower().split())] = r["assistant"]
    rows = [{"user": k, "assistant": v} for k, v in dedup.items()]

    out = os.path.join(C.DATA_ROOT, "sft_raw", "math_big.jsonl")
    with open(out, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(f"算术数据 -> {out}")
    print(f"  不同表达式 {len(items)} 个 -> {len(rows)} 条问答")
    print(f"  覆盖：加法 {sum(1 for e,_ in items if '+' in e)}，"
          f"减法 {sum(1 for e,_ in items if '-' in e)}，"
          f"乘法 {sum(1 for e,_ in items if '*' in e)}，"
          f"除法 {sum(1 for e,_ in items if '/' in e)}")
    print("\n样例：")
    for r in rows[:5] + rows[len(rows)//2:len(rows)//2+3]:
        print(f"  Q: {r['user']:<26} A: {r['assistant']}")
    # 关键校验：答案里的数字是否只出现在题目里（那说明模型能抄）
    import re as _re
    def nums(s):
        return set(_re.findall(r"\d+", s))
    copyable = sum(1 for r in rows
                   if nums(r["assistant"]) and nums(r["assistant"]) <= nums(r["user"]))
    print(f"\n  答案数字全都能在题目里找到的样本：{copyable}/{len(rows)}"
          f"（比例越低，说明模型越没法靠抄数字蒙对）")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())