"""Pinned line-level OCR protocol. Local Transformers classes; no remote Python."""
from __future__ import annotations
import hashlib, importlib.metadata, json, platform, sys, time
from pathlib import Path

def load_rows(task):
    """Read the pinned parquet locally, without a streaming HTTP worker at exit."""
    import io
    import pyarrow.parquet as pq
    from huggingface_hub import hf_hub_download
    from PIL import Image
    file=hf_hub_download(task['dataset'],f"data/{task['split']}.parquet",repo_type='dataset',revision=task['dataset_revision'])
    with pq.ParquetFile(file) as parquet:
        table=parquet.read(columns=['image','text'],use_threads=False)
        records=table.to_pylist()
    del table
    selected=select_rows(records,task['limit'])
    for _,row,_ in selected:
        blob=row['image']
        if not isinstance(blob,dict) or not isinstance(blob.get('bytes'),bytes):raise ValueError('Pinned OCR image bytes missing')
        with Image.open(io.BytesIO(blob['bytes'])) as image:row['image']=image.convert('RGB')
    return selected

def clean_text(value):
    """Preserve case, punctuation and internal whitespace; strip only the edges."""
    if not isinstance(value,str):raise ValueError('OCR transcription must be text')
    return value.strip()

def select_rows(dataset,limit):
    if not isinstance(limit,int) or isinstance(limit,bool) or not 1<=limit<=2915:raise ValueError('Invalid OCR sample limit')
    rows=[]
    for index,row in enumerate(dataset):
        reference=clean_text(row['text'])
        if not reference:continue
        rows.append((index,row,reference))
        if len(rows)==limit:return rows
    raise ValueError('Dataset did not yield the declared number of nonempty transcriptions')

def evaluate(job,output):
    import numpy as np
    import torch
    from jiwer import cer
    from transformers import TrOCRProcessor,VisionEncoderDecoderModel
    task=job['task'];model_id=job['model'];revision=job['revision']
    if task['adapter']!='ocr' or task['metric']!='cer' or task['direction']!='min':raise ValueError('Unsupported OCR protocol')
    if not any(m['id']==model_id and m['revision']==revision for m in task['models']):raise ValueError('Model is not pinned in this protocol')
    torch.manual_seed(42);np.random.seed(42);torch.set_num_threads(4)
    torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
    device='cuda:0' if torch.cuda.is_available() else 'cpu'
    started=time.monotonic()
    rows=load_rows(task);data_seconds=time.monotonic()-started
    loaded=time.monotonic()
    processor=TrOCRProcessor.from_pretrained(model_id,revision=revision,trust_remote_code=False,use_fast=False)
    model=VisionEncoderDecoderModel.from_pretrained(model_id,revision=revision,trust_remote_code=False,torch_dtype=torch.float32).to(device).eval()
    parameters=sum(p.numel() for p in model.parameters());load_seconds=time.monotonic()-loaded
    samples=[];scoring=time.monotonic()
    for index,row,reference in rows:
        began=time.monotonic();image=row['image'].convert('RGB')
        pixels=processor(images=image,return_tensors='pt').pixel_values.to(device)
        with torch.inference_mode():
            tokens=model.generate(pixels,max_new_tokens=128,do_sample=False,num_beams=1)
        prediction=clean_text(processor.batch_decode(tokens,skip_special_tokens=True)[0])
        samples.append(dict(id=str(index),reference=reference,prediction=prediction,score=float(cer(reference,prediction)),reference_characters=len(reference),image_size=list(image.size),image_rgb_sha256=hashlib.sha256(image.tobytes()).hexdigest(),seconds=time.monotonic()-began))
    scoring_seconds=time.monotonic()-scoring
    # Corpus CER weights each edit by total reference characters; it is not mean line CER.
    score=float(cer([s['reference'] for s in samples],[s['prediction'] for s in samples]))
    result=dict(score=score,n=len(samples),model_revision=revision,dataset_revision=task['dataset_revision'],runner_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),seconds=time.monotonic()-started,data_loading_seconds=data_seconds,model_load_seconds=load_seconds,scoring_seconds=scoring_seconds,peak_gpu_bytes=torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0,selection=f"First {task['limit']} nonempty transcriptions in pinned {task['split']} split order",seed=42,device=device,protocol=task,parameters=parameters,dtype='float32',tf32=False,batch_size=1,generation=dict(max_new_tokens=128,do_sample=False,num_beams=1),normalization='Strip outer whitespace only; preserve case, punctuation and internal whitespace',aggregation='Total character edits / total reference characters (corpus CER)',reference_characters=sum(s['reference_characters'] for s in samples),measured_cost_usd=None,notes='256-line subset of 2,915 IAM test lines. English handwriting only; not a document OCR/layout/table benchmark or official full IAM score.')
    payload=dict(result=result,samples=samples,environment=dict(python=sys.version,platform=platform.platform(),torch=torch.__version__,cuda=torch.version.cuda,gpu=torch.cuda.get_device_name() if torch.cuda.is_available() else None,packages={d.metadata['Name']:d.version for d in importlib.metadata.distributions()}))
    output.parent.mkdir(parents=True,exist_ok=True);output.write_text(json.dumps(payload,allow_nan=False))

if __name__=='__main__':evaluate(json.loads(Path(sys.argv[1]).read_text()),Path(sys.argv[2]))