"""Reuse fixed RoBERTa evaluation splits; prepare emotion rows for Gemma."""
import hashlib
import json
import os
from pathlib import Path
ROOT=Path(__file__).resolve().parent
os.environ.setdefault('HF_HOME',str(ROOT/'.cache/huggingface'))
from datasets import load_dataset


def main():
    out=ROOT/'data/open-emotions-gemma'
    if out.exists(): raise RuntimeError('Refusing to overwrite Gemma emotion data')
    out.mkdir()
    reference=ROOT/'outputs/open-emotions/run.json'
    if not reference.exists():reference=ROOT/'roberta_reference.json'
    old=json.loads(reference.read_text())
    names=old['labels'];ds=load_dataset(old['dataset'],'simplified',revision=old['dataset_revision'])
    sizes={}
    for split in ['train','validation','test']:
        rows=[]
        for row in ds[split]:
            rows.append(dict(id=row['id'],text=row['text'],targets=[float(i in row['labels']) for i in range(len(names))]))
        (out/f'{split}.jsonl').write_text(''.join(json.dumps(r,ensure_ascii=False)+'\n' for r in rows),encoding='utf-8')
        sizes[split]=len(rows)
    heldout={r['text'].casefold().strip() for split in ['validation','test'] for r in ds[split]}
    generated=[];removed=[]
    for line in (ROOT/'data/synthetic-emotions-v1.jsonl').read_text(encoding='utf-8').splitlines():
        row=json.loads(line)
        if row['text'].casefold().strip() in heldout: removed.append(row['id']);continue
        generated.append({**row,'targets':[float(n in row['labels']) for n in names]})
    (out/'synthetic.jsonl').write_text(''.join(json.dumps(r,ensure_ascii=False)+'\n' for r in generated),encoding='utf-8')
    meta=dict(dataset=old['dataset'],revision=old['dataset_revision'],labels=names,sizes=sizes,
              source_split='Same official GoEmotions simplified train/validation/test as RoBERTa baseline',
              synthetic_examples=len(generated),excluded_exact_heldout_matches=removed,
              synthetic_limitations='Previously assistant-authored 300 framing variants from 60 seeds focused on six rare emotions; no independent human review.',
              license='apache-2.0',input_format='plain text',
              files_sha256={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in out.iterdir()})
    (out/'data_manifest.json').write_text(json.dumps(meta,indent=2)+'\n',encoding='utf-8')
    print('Emotion data prepared',sizes,'synthetic',len(generated))


if __name__=='__main__':main()