#!/usr/bin/env python3
"""Enable a project-local custom-only SFT mode in nanochat chat_sft.py."""
from __future__ import annotations

from pathlib import Path


TARGET = Path("/home/bbabajan/nanochat-ar/nanochat/scripts/chat_sft.py")


OLD = '''# SFT data mixture and DataLoader
identity_conversations_filepath = os.path.join(base_dir, "identity_conversations.jsonl")
train_tasks = [
    SmolTalk(split="train"), # 460K rows of general conversations
    CustomJSON(filepath=identity_conversations_filepath), # 1000 rows of synthetic identity conversations
    CustomJSON(filepath=identity_conversations_filepath), # 2 epochs of these
    *[MMLU(subset="all", split="auxiliary_train") for _ in range(args.mmlu_epochs)], # 100K rows per epoch
    *[GSM8K(subset="main", split="train") for _ in range(args.gsm8k_epochs)], # 8K rows per epoch
    SimpleSpelling(size=200000, split="train"), # 200K rows of Simple Spelling (e.g. spell the word 'apple')
    SpellingBee(size=80000, split="train"), # 80K rows of Spelling Bee (e.g. how many 'r' are in 'strawberry'?)
]
train_dataset = TaskMixture(train_tasks)
print0(f"Training mixture: {len(train_dataset):,} rows (MMLU x{args.mmlu_epochs}, GSM8K x{args.gsm8k_epochs})")
val_dataset = TaskMixture([
    SmolTalk(split="test"), # 24K rows in test set
    MMLU(subset="all", split="test", stop=5200), # 14K rows in test set, use only 5.2K to match the train ratios
    GSM8K(subset="main", split="test", stop=420), # 1.32K rows in test set, use only 420 to match the train ratios
]) # total: 24K + 5.2K + 0.42K ~= 29.6K rows
'''


NEW = '''# SFT data mixture and DataLoader
custom_sft_path = os.environ.get("PROJECT7_CUSTOM_SFT_JSONL")
if custom_sft_path:
    train_dataset = TaskMixture([CustomJSON(filepath=custom_sft_path)])
    val_dataset = TaskMixture([CustomJSON(filepath=custom_sft_path)])
    print0(f"Training mixture: {len(train_dataset):,} rows from PROJECT7_CUSTOM_SFT_JSONL={custom_sft_path}")
else:
    identity_conversations_filepath = os.path.join(base_dir, "identity_conversations.jsonl")
    train_tasks = [
        SmolTalk(split="train"), # 460K rows of general conversations
        CustomJSON(filepath=identity_conversations_filepath), # 1000 rows of synthetic identity conversations
        CustomJSON(filepath=identity_conversations_filepath), # 2 epochs of these
        *[MMLU(subset="all", split="auxiliary_train") for _ in range(args.mmlu_epochs)], # 100K rows per epoch
        *[GSM8K(subset="main", split="train") for _ in range(args.gsm8k_epochs)], # 8K rows per epoch
        SimpleSpelling(size=200000, split="train"), # 200K rows of Simple Spelling (e.g. spell the word 'apple')
        SpellingBee(size=80000, split="train"), # 80K rows of Spelling Bee (e.g. how many 'r' are in 'strawberry'?)
    ]
    train_dataset = TaskMixture(train_tasks)
    print0(f"Training mixture: {len(train_dataset):,} rows (MMLU x{args.mmlu_epochs}, GSM8K x{args.gsm8k_epochs})")
    val_dataset = TaskMixture([
        SmolTalk(split="test"), # 24K rows in test set
        MMLU(subset="all", split="test", stop=5200), # 14K rows in test set, use only 5.2K to match the train ratios
        GSM8K(subset="main", split="test", stop=420), # 1.32K rows in test set, use only 420 to match the train ratios
    ]) # total: 24K + 5.2K + 0.42K ~= 29.6K rows
'''


def main() -> int:
    text = TARGET.read_text(encoding="utf-8")
    if "PROJECT7_CUSTOM_SFT_JSONL" in text:
        print("chat_sft.py already supports PROJECT7_CUSTOM_SFT_JSONL")
        return 0
    if OLD not in text:
        raise SystemExit("Could not find expected chat_sft.py SFT mixture block")
    TARGET.write_text(text.replace(OLD, NEW), encoding="utf-8")
    print("Patched chat_sft.py for PROJECT7_CUSTOM_SFT_JSONL")
    return 0


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