"""
Build the Japanese-heavy conversational calibration set
(calib_conversational.jsonl). An English-only UltraChat draw of the same
size is also written as a byproduct (calib_en_ultrachat.jsonl).

Mix (Section 6 of the project plan), target 512 samples total:
  - 40% JA multi-turn instruction/chat  -> llm-jp/oasst2-33k-ja
  - 20% JA knowledge/long-form          -> JA Wikipedia passages, instruction-wrapped
  - 15% JA reasoning/math               -> Kendamarron/magpie-japanese-math-instruction-17k
  - 10% Code (JA-instructed)            -> ronantakizawa/python-code-instructions-japanese
  - 15% English chat                    -> HuggingFaceH4/ultrachat_200k

Every sample is stored as a `messages` list (chat format) and later fed
through `tokenizer.apply_chat_template(messages, add_generation_prompt=False)`
at PTQ time, so calibration activations see the same role tokens /
`<|...|>` specials / structure the model sees in real deployment.

IMPORTANT: the Wikipedia articles used here for the knowledge/long-form slice
must stay disjoint from the articles used in perplexity_ja.py's held-out set
(recorded in results/*/perplexity_ja.json `corpus.article_ids`). The held-out
set consumed articles from the very start of the stream (skip 0, ~50 articles);
we skip --wiki-skip-articles (default 5000) before collecting, guaranteeing
disjointness with two orders of magnitude of margin.

The EN-only set is drawn from the SAME single UltraChat pass as the
main mix's en_chat slice and split by index, so the two are disjoint by
construction (no skip heuristics).
"""

import argparse
import json
import random

from datasets import load_dataset
from transformers import AutoTokenizer

SEED = 1234
TARGET_TOTAL = 512
MAX_TOKENS = 4096

MIX = {
    "ja_chat": 0.40,
    "ja_knowledge": 0.20,
    "ja_math": 0.15,
    "code": 0.10,
    "en_chat": 0.15,
}


def iter_ja_chat(rng):
    # oasst2-33k-ja alone fills the ja_chat pool target, so it is the sole
    # source actually drawn from (see the datasheet).
    sources = [
        ("llm-jp/oasst2-33k-ja", "apache-2.0"),
    ]
    for name, license_ in sources:
        ds = load_dataset(name, split="train", streaming=True)
        for ex in ds:
            convo = ex["conversations"]
            if len(convo) < 2:
                continue
            yield {
                "messages": [{"role": m["role"], "content": m["content"]} for m in convo],
                "source": name,
                "license": license_,
            }


def iter_ja_knowledge(skip_articles: int):
    ds = load_dataset("wikimedia/wikipedia", "20231101.ja", split="train", streaming=True)
    it = iter(ds)
    for _ in range(skip_articles):
        next(it)
    prompts = [
        "次のトピックについて、知っていることを詳しく説明してください：{title}",
        "「{title}」について解説してください。",
        "次の見出し語を要約・解説する文章を書いてください：{title}",
    ]
    i = 0
    for ex in it:
        text = ex["text"].strip()
        if len(text) < 800:
            continue
        prompt = prompts[i % len(prompts)].format(title=ex["title"])
        yield {
            "messages": [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": text[:3000]},
            ],
            "source": f"wikimedia/wikipedia:20231101.ja#{ex['id']}",
            "license": "cc-by-sa-4.0",
        }
        i += 1


def iter_ja_math():
    ds = load_dataset(
        "Kendamarron/magpie-japanese-math-instruction-17k-qwen2.5-bakeneko-32b-instruct",
        split="train",
        streaming=True,
    )
    for ex in ds:
        if not ex.get("instruction") or not ex.get("output"):
            continue
        yield {
            "messages": [
                {"role": "user", "content": ex["instruction"]},
                {"role": "assistant", "content": ex["output"]},
            ],
            "source": "Kendamarron/magpie-japanese-math-instruction-17k",
            "license": "apache-2.0",
        }


def iter_code():
    ds = load_dataset("ronantakizawa/python-code-instructions-japanese", split="train", streaming=True)
    for ex in ds:
        instruction = ex.get("instruction", "")
        input_ = ex.get("input", "")
        output = ex.get("output", "")
        if not instruction or not output:
            continue
        user_content = instruction if not input_ else f"{instruction}\n\n{input_}"
        yield {
            "messages": [
                {"role": "user", "content": user_content},
                {"role": "assistant", "content": output},
            ],
            "source": "ronantakizawa/python-code-instructions-japanese",
            "license": "mit",
        }


def iter_en_chat():
    ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft", streaming=True)
    for ex in ds:
        messages = ex.get("messages", [])
        if len(messages) < 2:
            continue
        yield {
            "messages": [{"role": m["role"], "content": m["content"]} for m in messages],
            "source": "HuggingFaceH4/ultrachat_200k",
            "license": "mit",
        }


def collect(gen, tokenizer, count, rng, max_tokens=MAX_TOKENS, pool_multiplier=4, label=""):
    """Pull samples from a generator, keep ones that fit the token budget,
    then randomly downsample to `count` for an unbiased draw from the pool."""
    pool = []
    pool_target = count * pool_multiplier
    for sample in gen:
        # Template to a string, then encode — apply_chat_template(tokenize=True)'s
        # return type varies across transformers versions (list of ids vs
        # BatchEncoding, whose len() is its dict key count and silently
        # broke this filter on 5.5.x). This two-step form is unambiguous.
        templated = tokenizer.apply_chat_template(
            sample["messages"], add_generation_prompt=False, tokenize=False
        )
        n_tokens = len(tokenizer.encode(templated, add_special_tokens=False))
        if n_tokens > max_tokens or n_tokens < 8:
            continue
        sample["num_tokens"] = n_tokens
        pool.append(sample)
        if len(pool) % 100 == 0:
            print(f"  [{label}] pool {len(pool)}/{pool_target}", flush=True)
        if len(pool) >= pool_target:
            break
    rng.shuffle(pool)
    return pool[:count]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--tokenizer", default="./models/llm-jp-4-8b-instruct")
    parser.add_argument("--output-dir", default="data/calib")
    parser.add_argument("--total", type=int, default=TARGET_TOTAL)
    parser.add_argument(
        "--wiki-skip-articles",
        type=int,
        default=5000,
        help="Skip this many leading JA-Wikipedia articles to stay well clear "
        "of perplexity_ja.py's held-out set (which consumes articles from "
        "the very start of the stream).",
    )
    args = parser.parse_args()

    rng = random.Random(SEED)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)

    counts = {k: round(v * args.total) for k, v in MIX.items()}
    # Fix up rounding so counts sum exactly to args.total.
    diff = args.total - sum(counts.values())
    counts[max(counts, key=counts.get)] += diff
    print("Target mix:", counts)

    all_samples = []
    datasheet = {"seed": SEED, "total": args.total, "max_tokens": MAX_TOKENS, "mix": {}}

    print("Collecting ja_chat...", flush=True)
    ja_chat = collect(iter_ja_chat(rng), tokenizer, counts["ja_chat"], rng, label="ja_chat")
    all_samples += ja_chat

    print("Collecting ja_knowledge...", flush=True)
    ja_knowledge = collect(
        iter_ja_knowledge(args.wiki_skip_articles),
        tokenizer,
        counts["ja_knowledge"],
        rng,
        label="ja_knowledge",
    )
    all_samples += ja_knowledge

    print("Collecting ja_math...", flush=True)
    ja_math = collect(iter_ja_math(), tokenizer, counts["ja_math"], rng, label="ja_math")
    all_samples += ja_math

    print("Collecting code...", flush=True)
    code = collect(iter_code(), tokenizer, counts["code"], rng, label="code")
    all_samples += code

    # Collect both EN slices in one filtered pass over one UltraChat
    # stream, then split by index — disjoint by construction.
    print("Collecting en_chat (single pass)...", flush=True)
    en_all = collect(
        iter_en_chat(),
        tokenizer,
        counts["en_chat"] + args.total,
        rng,
        pool_multiplier=2,
        label="en_chat",
    )
    en_chat = en_all[: counts["en_chat"]]
    en_only_extra = en_all[counts["en_chat"] :]
    all_samples += en_chat

    for name, samples in [
        ("ja_chat", ja_chat),
        ("ja_knowledge", ja_knowledge),
        ("ja_math", ja_math),
        ("code", code),
        ("en_chat", en_chat),
    ]:
        sources = sorted(set(s["source"].split("#")[0] for s in samples))
        licenses = sorted(set(s["license"] for s in samples))
        datasheet["mix"][name] = {
            "count": len(samples),
            "target": counts[name],
            "sources": sources,
            "licenses": licenses,
            "mean_tokens": sum(s["num_tokens"] for s in samples) / max(len(samples), 1),
        }

    datasheet["en_only_extra"] = {
        "count": len(en_only_extra),
        "target": args.total,
        "sources": ["HuggingFaceH4/ultrachat_200k"],
        "licenses": ["mit"],
        "mean_tokens": sum(s["num_tokens"] for s in en_only_extra) / max(len(en_only_extra), 1),
        "note": "drawn from the same single-pass pool as the main mix's "
        "en_chat slice and split by index — disjoint by construction",
    }

    rng.shuffle(all_samples)

    out_path = f"{args.output_dir}/calib_conversational.jsonl"
    with open(out_path, "w") as f:
        for s in all_samples:
            f.write(json.dumps({"messages": s["messages"], "source": s["source"]}, ensure_ascii=False) + "\n")

    datasheet["actual_total"] = len(all_samples)
    with open(f"{args.output_dir}/calibration_datasheet_conversational.json", "w") as f:
        json.dump(datasheet, f, ensure_ascii=False, indent=2)

    print(f"\nWrote {len(all_samples)} samples to {out_path}", flush=True)
    print(json.dumps(datasheet, ensure_ascii=False, indent=2), flush=True)

    en_out_path = f"{args.output_dir}/calib_en_ultrachat.jsonl"
    with open(en_out_path, "w") as f:
        for s in en_only_extra:
            f.write(json.dumps({"messages": s["messages"], "source": s["source"]}, ensure_ascii=False) + "\n")
    print(f"Wrote {len(en_only_extra)} EN-only samples to {en_out_path}", flush=True)


if __name__ == "__main__":
    main()