#!/usr/bin/env python3
"""Run from project root: .venv/bin/python scripts/prepare_data.py."""
import sys
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import json, math, hashlib, bisect
from collections import Counter, defaultdict
import numpy as np
import pyarrow.parquet as pq
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer
from pii.data import LABELS, DATA_REVISION, MODEL_REVISION, validate, digest, family_key, support

ROOT=Path(__file__).resolve().parents[1]
RAW=ROOT/'data/raw'; OUT=ROOT/'data/prepared'; ART=ROOT/'artifacts'
OUT.mkdir(parents=True,exist_ok=True); ART.mkdir(exist_ok=True)
manifest={'source':'nvidia/Nemotron-PII','revision':DATA_REVISION,'license':'CC-BY-4.0','attribution':'NVIDIA; Amy Steier, Andre Manoel, Alexa Haushalter, Maarten Van Segbroeck, Nemotron-PII (2025). https://huggingface.co/datasets/nvidia/Nemotron-PII','source_files':{},'source_stats':{},'exclusions':{}}
files=['README.md']+[f'data/{s}-00000-of-00001.parquet' for s in ('train','test')]
for f in files:
    p=Path(hf_hub_download(manifest['source'],f,repo_type='dataset',revision=DATA_REVISION,local_dir=RAW))
    manifest['source_files'][f]={'sha256':hashlib.sha256(p.read_bytes()).hexdigest(),'bytes':p.stat().st_size}
rows=[]; excluded=[]; parent=[]; key_owner={}
def find(i):
    while parent[i]!=i:
        parent[i]=parent[parent[i]]; i=parent[i]
    return i
def union(i,j):
    a,b=find(i),find(j)
    if a!=b: parent[max(a,b)]=min(a,b)
for source in ('train','test'):
    table=pq.read_table(RAW/f'data/{source}-00000-of-00001.parquet')
    stats={'rows':len(table),'labels':Counter(),'invalid':Counter(),'empty_all_spans':0,'semantic_flags':Counter()}
    seen_uids=Counter(); seen_texts=Counter()
    for n,r in enumerate(table.to_pylist()):
        seen_uids[r['uid']]+=1; seen_texts[digest(r['text'])]+=1
        spans,reason=validate(r)
        if reason:
            stats['invalid'][reason]+=1; excluded.append({'source':source,'index':n,'uid':r['uid'],'reason':reason}); continue
        stats['labels'].update(s['label'] for s in spans); stats['empty_all_spans']+=not spans
        for s in spans:
            if s['label']=='email' and '@' not in s['text']:
                stats['semantic_flags']['email_without_at']+=1
        row={'id':f'{source}:{n}:{r["uid"]}','uid':r['uid'],'source':source,'source_index':n,'text':r['text'],'spans':[{k:s[k] for k in ('label','start','end')} for s in spans if s['label'] in LABELS], 'domain':r['domain'],'document_type':r['document_type'],'family':family_key(r['text'],spans),'text_sha256':digest(r['text'])}
        i=len(rows); parent.append(i)
        for key in ('uid:'+row['uid'],'text:'+row['text_sha256'],'family:'+row['family']):
            if key in key_owner: union(i,key_owner[key])
            else:key_owner[key]=i
        rows.append(row)
    stats['duplicate_uid_rows']=sum(v-1 for v in seen_uids.values()); stats['duplicate_exact_text_rows']=sum(v-1 for v in seen_texts.values())
    manifest['source_stats'][source]=stats
    print(source,json.dumps(stats),flush=True)
groups=defaultdict(list)
for i,row in enumerate(rows): groups[find(i)].append(row)
splits={s:[] for s in ('train','validation','calibration_fit','calibration_select','test')}
train_groups=[]; overlaps=0
for group in groups.values():
    group_id=min(r['id'] for r in group)
    for r in group:r['group']=digest(group_id)
    test=[r for r in group if r['source']=='test']
    if test:
        splits['test'].extend(test)
        for r in group:
            if r['source']=='train':
                excluded.append({'source':'train','index':r['source_index'],'uid':r['uid'],'reason':'test_family_overlap'}); overlaps+=1
    else:train_groups.append(group)
# Stable pseudorandom family order. Allocate heldouts jointly, without model predictions.
train_groups.sort(key=lambda g:digest('datafog-pii-split-v1:'+g[0]['group']))
# Owner-approved transparent volume-derived heuristic; no quality guarantee.
all_support=Counter()
for group in train_groups:
    for r in group: all_support.update(set(s['label'] for s in r['spans']) or ['__negative__'])
quota=math.ceil(math.sqrt(min(all_support[l] for l in LABELS)))
heldouts=('validation','calibration_fit','calibration_select'); counts={s:Counter() for s in heldouts}
for group in train_groups:
    coverage=Counter()
    for r in group:
        labels=set(s['label'] for s in r['spans']); coverage.update(labels or ['__negative__'])
    deficits={s:sum(min(v,max(0,quota-counts[s][l])) for l,v in coverage.items()) for s in heldouts}
    chosen=max(heldouts,key=lambda s:(deficits[s],-len(splits[s])))
    if deficits[chosen]==0:chosen='train'
    splits[chosen].extend(group)
    if chosen in counts:counts[chosen].update(coverage)
manifest['split_rule']={'quota_positive_documents_per_label_and_negative_documents':quota,'derivation':'ceil(sqrt(minimum observed eligible positive-document support across target labels)); equal support in three heldouts. Transparent prototype allocation heuristic balancing training retention with heldout growth; no statistical precision or quality guarantee.','eligible_train_document_support':dict(all_support),'grouping':'Connected components of UID, exact original text SHA256, and whitespace-collapsed casefolded all-label-masked text SHA256. All train groups touching test excluded. Source test retained in full except structurally invalid rows.','allocation':'Stable SHA256 family order; assign each family to heldout with greatest remaining document-support deficit (tie: smaller partition), otherwise train.','seed':'datafog-pii-split-v1','test_overlap_train_rows':overlaps,'heldout_achieved_support':counts}
manifest['splits']={}
for name,records in splits.items():
    records.sort(key=lambda r:r['id'])
    p=OUT/f'{name}.jsonl'
    with p.open('w') as f:
        for r in records:f.write(json.dumps(r,ensure_ascii=False)+'\n')
    manifest['splits'][name]={**support(records),'sha256':hashlib.sha256(p.read_bytes()).hexdigest()}
    print(name,json.dumps(manifest['splits'][name]),flush=True)
sets={name:set(r['group'] for r in records) for name,records in splits.items()}
for a in sets:
    for b in sets:
        if a<b:assert sets[a].isdisjoint(sets[b]),(a,b)
manifest['pairwise_group_disjoint']=True
manifest['exclusions']=dict(Counter(r['reason'] for r in excluded))
with (OUT/'excluded.jsonl').open('w') as f:
    for r in excluded:f.write(json.dumps(r)+'\n')
(ART/'data-manifest.json').write_text(json.dumps(manifest,indent=2))
print('Tokenizing full eligible training-side data for length/alignment statistics',flush=True)
tokenizer=AutoTokenizer.from_pretrained('distilbert/distilbert-base-cased',revision=MODEL_REVISION)
tokenizer.save_pretrained(ROOT/'artifacts/tokenizer')
length_stats={}; entity_lengths=defaultdict(list); boundary_bad=Counter(); boundary_examples=[]
for name in ('train','validation','calibration_fit','calibration_select'):
    lengths=[]
    for base in range(0,len(splits[name]),256):
        batch=splits[name][base:base+256]
        encoded=tokenizer([r['text'] for r in batch],add_special_tokens=False,return_offsets_mapping=True,truncation=False,verbose=False)
        for r,ids,offsets in zip(batch,encoded['input_ids'],encoded['offset_mapping']):
            lengths.append(len(ids)); starts=[a for a,b in offsets]; ends=[b for a,b in offsets]
            for s in r['spans']:
                i=bisect.bisect_right(ends,s['start']); j=bisect.bisect_left(starts,s['end'])
                entity_lengths[s['label']].append(j-i)
                if i>=j or starts[i]!=s['start'] or ends[j-1]!=s['end']:
                    boundary_bad[s['label']]+=1
                    if len(boundary_examples)<30:boundary_examples.append({'id':r['id'],'span':s,'text':r['text'][s['start']:s['end']],'covering_offsets':offsets[i:j]})
    length_stats[name]={'documents':len(lengths),'tokens':sum(lengths),'max':max(lengths,default=0),'quantiles':dict(zip(['p50','p90','p95','p99'],np.quantile(lengths,[.5,.9,.95,.99]).tolist())),'over_510_tokens':sum(x>510 for x in lengths),'characters':sum(len(r['text']) for r in splits[name])}
    print(name,length_stats[name],flush=True)
manifest['tokenizer']={'model':'distilbert/distilbert-base-cased','revision':MODEL_REVISION,'model_max_length':tokenizer.model_max_length,'lengths':length_stats,'entity_tokens_by_label':{k:{'count':len(v),'max':max(v),'p99':float(np.quantile(v,.99))} for k,v in entity_lengths.items()},'unrepresentable_boundaries_by_label':dict(boundary_bad),'unrepresentable_boundary_examples':boundary_examples}
(ART/'data-manifest.json').write_text(json.dumps(manifest,indent=2))
print('DONE',json.dumps(manifest['tokenizer']),flush=True)