import hashlib
import json
import re
from common import ROOT, rows, save_json


def dump(path, items):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(''.join(json.dumps(r, ensure_ascii=False) + '\n' for r in items))


def main():
    groups = {}
    for name in ('gsm8k_train', 'math_train', 'aime'):
        for row in rows(ROOT / f'data/practice/{name}.jsonl'):
            canonical = ' '.join(re.findall(r'\w+', row['problem'].lower()))
            digest = hashlib.sha256(('sol-independent-v1:' + canonical).encode()).digest()
            bucket = int.from_bytes(digest[:8], 'big') % 10
            split = 'discovery' if bucket < 7 else 'validation' if bucket < 9 else 'confirmation'
            groups.setdefault(split, []).append(row)
    for split, records in groups.items():
        dump(ROOT / f'data/splits/{split}.jsonl', records)
    diagnostic = []
    for source, count in [('gsm8k', 32), ('math', 64)]:
        subset = [r for r in groups['discovery'] if r['source'] == source]
        diagnostic.extend(sorted(subset, key=lambda r: hashlib.sha256(r['id'].encode()).hexdigest())[:count])
    dump(ROOT / 'data/diagnostic.jsonl', diagnostic)
    manifest = {k: {'count': len(v), 'sha256': hashlib.sha256((ROOT / f'data/splits/{k}.jsonl').read_bytes()).hexdigest()} for k, v in groups.items()}
    manifest['_method'] = 'sha256(sol-independent-v1:normalized_problem), 70/20/10 buckets; identical normalized texts stay together'
    save_json(ROOT / 'data/splits/manifest.json', manifest)
    print(json.dumps({k: len(v) for k, v in groups.items()}))


if __name__ == '__main__':
    main()