"""Validate Gemma 4's real tokenizer against window IDs and evidence spans."""
import os
import time
from pathlib import Path
os.environ.setdefault('HF_HOME',str(Path(__file__).resolve().parent/'.cache/huggingface'))
from transformers import AutoTokenizer
from gemma_classifier import BASE,REVISION
from predict_windows import make_windows

def main():
    tokenizer=AutoTokenizer.from_pretrained(BASE,revision=REVISION)
    short='The aroma of café coffee filled the kitchen.'
    windows,_=make_windows(tokenizer,short,64,0)
    assert windows[0]['encoded']['input_ids']==tokenizer(short)['input_ids']
    text=(short+' ')*60
    for target,overlap in [(None,0),(None,16),('aroma',16)]:
        prefix=f'Target: {target}\nContext: ' if target else ''
        original=tokenizer(prefix+text,add_special_tokens=False,return_offsets_mapping=True)
        count=sum(end<=len(prefix) for _,end in original['offset_mapping']) if prefix else 0
        prefix_ids=original['input_ids'][:count];context_ids=original['input_ids'][count:]
        windows,metadata=make_windows(tokenizer,text,64,overlap,target)
        assert len(windows)>1
        for window in windows:
            assert window['encoded']['input_ids']==[tokenizer.bos_token_id]+prefix_ids+context_ids[window['token_start']:window['token_end']]
            assert window['input_tokens']<=64
            assert text[window['char_start']:window['char_end']]==window['text']
            if target:assert target in window['text']
        print('Exact Gemma token IDs/spans passed',target or 'text','overlap',overlap,'windows',len(windows),flush=True)
    windows,metadata=make_windows(tokenizer,short+' '+'An ordinary distant event. '*100,64,16,'aroma')
    assert metadata['skipped_windows_without_complete_target']>0
    assert all('aroma' in window['text'] for window in windows)
    print('Target omission guard passed',flush=True)
    long_text=(short+' ')*400
    assert len(long_text)<=20000
    started=time.monotonic()
    windows,_=make_windows(tokenizer,long_text,128,32,'aroma')
    assert 20<len(windows)<=128 and all(w['input_tokens']<=128 for w in windows)
    print('Long sensory document passed',len(long_text),'characters',len(windows),
          'windows',round(time.monotonic()-started,3),'seconds',flush=True)

if __name__=='__main__':main()