"""
Push the synthetic Arabic math-reasoning corpus to the Hub as a PRIVATE dataset with a card.

Usage: python push_synth_dataset.py [--repo oddadmix/arabic-math-reasoning-synth] [--dry-run]
"""
import argparse
import collections
import json
import os
from pathlib import Path

import pyarrow.parquet as pq
from huggingface_hub import HfApi

OUT_DIR = Path(os.environ.get("OUT_DIR", "out_synth"))
PARQUET = OUT_DIR / "arabic_math_reasoning_synth.parquet"

# Hub ids for the generators, keyed by the `model` field the nodes record per row.
GEN_MODEL_HUB = {
    "gemma-3-12b-it": "google/gemma-3-12b-it",
    "Qwen3.6-27B-NVFP4": "nvidia/Qwen3.6-27B-NVFP4",
    "Qwen3.8-27B-Uncensored-NVFP4": "orcarouter/Qwen3.8-27B-Uncensored-NVFP4",
}

CARD = """---
license: apache-2.0
language:
- ar
size_categories:
- {size_cat}
task_categories:
- text-generation
tags:
- arabic
- math
- reasoning
- chain-of-thought
- gsm8k
- synthetic
configs:
- config_name: default
  data_files:
  - split: train
    path: arabic_math_reasoning_synth.parquet
---

# Arabic Math Reasoning (synthetic) — مسائل رياضيات عربية مع خطوات الحل
{status_note}
**{n:,}** Arabic grade-school math word problems, each with a step-by-step derivation and a
concluding sentence. Generated with {gen_blurb} and
**arithmetically verified** — every equation the reasoning states was re-evaluated, and rows whose
own arithmetic does not check out were dropped.

{gen_table}

> **بالعربية:** {n:,} مسألة حسابية عربية مع خطوات حل مفصّلة وجملة إجابة نهائية. جميع المعادلات
> داخل خطوات الحل تم التحقق منها حسابيًا، والصفوف التي تحتوي على خطأ حسابي مستبعدة.

## Schema

| field | description |
|---|---|
| `instruction` | the word problem, ending in a question |
| `reasoning` | numbered derivation, each step containing an explicit equation |
| `answer` | one concluding sentence (`إذن، …`) carrying the final number and its unit |
| `axis_*` | the variation axes this item was drawn from (domain, operation, steps, numbers, region, twist) |
| `task_id` | the generation task it came from |
| `gen_model` | which model generated the row |

It matches the `{{instruction, reasoning, answer}}` schema used by the Nawah-Reasoning training
scripts, so it drops straight into ChatML + `<think>` SFT.

## How it was built

Prompts are drawn from a **variation grid** — {n_domains} domains × {n_ops} operation types ×
step counts × number styles × {n_regions} country/currency pairs × {n_twists} structural twists —
seeded deterministically from the task index. This is a deliberate reaction to the corpus this was
built to supplement: a machine-translated Arabic GSM8K whose {gsm_rows:,} rows expand from only
2,814 underlying question templates. Here, **{templates:,}** distinct templates survive
number-masked deduplication.
{pool_note}

### Verification

Generated math is fluent and often wrong, so acceptance is mechanical rather than stylistic:

1. every `a op b = c` in the reasoning is re-evaluated with real arithmetic — one bad equation
   rejects the row;
2. the number in the answer sentence must equal the last computed result (a chain that is
   internally correct but ends on a different number is the failure that survives fluency checks);
3. no Latin characters, length bounds, and the answer must open with a conclusion marker.

**{accept_rate:.1%}** of parsed items passed. Rejections:

{reject_table}

`out_synth/rejects.jsonl` in the source repo keeps every rejected item with its reason, so the
filter itself can be audited.

## Limitations

This is **synthetic** data: the problems are machine-written and reflect that model's habits of
phrasing, its distribution of scenarios, and its idea of what a "typical" Arabic math problem
looks like. Verification proves each row's arithmetic is self-consistent — it does **not** prove
the problem is well-posed, that the setup is the only reasonable reading, or that the wording is
natural to any particular dialect region (the `axis_region` field sets currency and context, not
dialect; everything is MSA). Rows whose reasoning states no explicit equation are dropped, which
biases the corpus toward problems that decompose into clean arithmetic steps.
"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="oddadmix/arabic-math-reasoning-synth")
    ap.add_argument("--target", type=int, default=100_000)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    import synth_common as sc

    n = pq.read_metadata(PARQUET).num_rows
    stats = json.loads((OUT_DIR / "build_stats.json").read_text(encoding="utf-8"))

    # Credit every generator that actually contributed, from the build's own per-model counts.
    by_model = collections.Counter({k[len("model_"):]: v
                                    for k, v in stats["stats"].items() if k.startswith("model_")})
    def _link(name):
        hub = GEN_MODEL_HUB.get(name)
        return f"[`{name}`](https://huggingface.co/{hub})" if hub else f"`{name}`"
    names = [m for m, _ in by_model.most_common()]
    gen_blurb = " and ".join(filter(None, [", ".join(_link(m) for m in names[:-1]),
                                           _link(names[-1])])) if names else "an unrecorded model"
    gen_table = ("| generator | rows | share |\n|---|---:|---:|\n" + "\n".join(
        f"| {_link(m)} | {c:,} | {100 * c / max(n, 1):.1f}% |"
        for m, c in by_model.most_common())) if len(names) > 1 else ""
    total_rejects = max(sum(stats["reject_reasons"].values()), 1)
    folded = collections.Counter()
    for k, v in stats["reject_reasons"].items():
        folded[k.split(":", 1)[0]] += v
    reject_table = "| reason | count | share of rejects |\n|---|---:|---:|\n" + "\n".join(
        f"| `{k}` | {v:,} | {100 * v / total_rejects:.1f}% |"
        for k, v in folded.most_common(8))

    if n < args.target:
        status_note = (
            f"\n> ⚠️ **Interim snapshot — generation is still running.** This is "
            f"{n:,} of a planned {args.target:,} rows, published early so it can be reviewed and "
            f"trained against. Rows are only ever *appended*: the generator is resumable and "
            f"deterministic per task index, so everything here stays in the finished corpus "
            f"unchanged. Expect this repo to be overwritten with a larger version.\n")
    else:
        status_note = ""

    # The relational pool (task ids from 1,000,000) is a second, disjoint operation set added
    # after the first 100,323 rows shipped. Count it from the data rather than asserting it.
    n_rel = sum(1 for tid in pq.read_table(PARQUET, columns=["task_id"]).column("task_id").to_pylist()
                if tid >= 1_000_000)
    pool_note = ("""
The operation axis is drawn from **two disjoint pools**. The first {n_default:,} rows use 15
general operations (totals, percentages, unit rates, remainders, …). A later pass added a
**relational pool** of 8 operations — *ضعف* / *أضعاف*, *نصف*, *ثلث* و *ربع*, absolute increase and
decrease, three-way chains, and the inverse direction — contributing **{n_rel:,}** rows. These
express a quantity *relative to another entity* ("Ziad bought **double that number**"), which
forces a derivation step that a percentage problem does not. In the original corpus *ضعف /
أضعاف* appeared in 1.34% of rows and *أكثر بـ / أقل بـ* in 0.26% — against 17.6% for percentages —
and that gap is the documented weak spot of models trained on it, so the pool was generated
deliberately rather than left to the grid.
""".format(n_rel=n_rel, n_default=n - n_rel) if n_rel else "")

    size_cat = "100K<n<1M" if n >= 100_000 else "10K<n<100K" if n >= 10_000 else "1K<n<10K"
    card = CARD.format(
        n=n, gen_blurb=gen_blurb, gen_table=gen_table, size_cat=size_cat,
        n_domains=len(sc.DOMAINS), n_ops=stats["axis_coverage"]["op"], n_regions=len(sc.REGIONS),
        n_twists=len(sc.TWISTS), gsm_rows=142_969,
        templates=stats["unique_templates"], accept_rate=stats["accept_rate"],
        status_note=status_note,
        reject_table=reject_table,
        pool_note=pool_note,
    )
    (OUT_DIR / "README.md").write_text(card, encoding="utf-8")
    print(f"[+] wrote card ({len(card)} chars) | {n:,} rows | accept {stats['accept_rate']:.1%}")

    if args.dry_run:
        print("[dry-run] not pushing")
        return

    api = HfApi()
    api.create_repo(args.repo, repo_type="dataset", private=True, exist_ok=True)
    api.upload_file(path_or_fileobj=str(PARQUET), path_in_repo=PARQUET.name,
                    repo_id=args.repo, repo_type="dataset")
    api.upload_file(path_or_fileobj=str(OUT_DIR / "README.md"), path_in_repo="README.md",
                    repo_id=args.repo, repo_type="dataset")
    for script in ("synth_common.py", "synth_generate.py", "build_synth_dataset.py"):
        api.upload_file(path_or_fileobj=script, path_in_repo=script,
                        repo_id=args.repo, repo_type="dataset")
    print(f"[+] https://huggingface.co/datasets/{args.repo}")


if __name__ == "__main__":
    main()