"""Save unfiltered predictions from the deployed ONNX model for evaluation."""
import argparse
from collections import Counter
import json
import sys
import time
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
from pii.core import decode_window, result, encode_windows

def main():
    p=argparse.ArgumentParser()
    p.add_argument("--bundle",type=Path,default=Path("artifacts/bundle"))
    p.add_argument("--input",type=Path,required=True)
    p.add_argument("--output",type=Path,required=True)
    p.add_argument("--batch-size",type=int,default=16,help="Measured training recipe batch; dynamically padded ONNX batches")
    args=p.parse_args()
    config=json.loads((args.bundle/"config.json").read_text())
    tokenizer=AutoTokenizer.from_pretrained(args.bundle,local_files_only=True)
    options=ort.SessionOptions();options.intra_op_num_threads=1
    session=ort.InferenceSession(str(args.bundle/"model.onnx"),sess_options=options,providers=["CPUExecutionProvider"])
    args.output.parent.mkdir(parents=True,exist_ok=True)
    started=time.perf_counter();documents=windows=0
    def batch_predict(rows,handle):
        nonlocal documents,windows
        encoded={"input_ids":[],"attention_mask":[],"offset_mapping":[]}
        mapping=[]
        for i,row in enumerate(rows):
            windows_for_row=encode_windows(tokenizer,row["text"],config["max_length"],config["stride"])
            for key in encoded:encoded[key].extend(windows_for_row[key])
            mapping.extend([i]*len(windows_for_row["input_ids"]))
        counts=Counter(mapping);seen=Counter();findings=[[] for _ in rows]
        for begin in range(0,len(mapping),args.batch_size):
            end=min(len(mapping),begin+args.batch_size)
            length=max(len(x) for x in encoded["input_ids"][begin:end])
            ids=np.full((end-begin,length),tokenizer.pad_token_id,dtype=np.int64)
            masks=np.zeros_like(ids)
            for j,index in enumerate(range(begin,end)):
                n=len(encoded["input_ids"][index]);ids[j,:n]=encoded["input_ids"][index];masks[j,:n]=encoded["attention_mask"][index]
            logits=session.run(None,{"input_ids":ids,"attention_mask":masks})[0]
            for j,index in enumerate(range(begin,end)):
                sample=mapping[index];offsets=encoded["offset_mapping"][index]
                findings[sample].extend(decode_window(logits[j,:len(offsets)],offsets,rows[sample]["text"],config["max_length"],has_previous=seen[sample]>0,has_next=seen[sample]+1<counts[sample]))
                seen[sample]+=1
        for row,found in zip(rows,findings):
            prediction={"id":row["id"],"text":row["text"],**result(row["text"],found)}
            handle.write(json.dumps(prediction,ensure_ascii=False)+"\n")
        documents+=len(rows);windows+=len(mapping)
        args.output.with_suffix(".progress.json").write_text(json.dumps({"documents":documents,"windows":windows,"seconds":time.perf_counter()-started})+"\n")
    temporary=args.output.with_suffix(".partial.jsonl")
    with args.input.open() as source,temporary.open("w") as target:
        rows=[]
        for line in source:
            if not line.strip():continue
            rows.append(json.loads(line))
            if len(rows)==args.batch_size:batch_predict(rows,target);rows=[]
        if rows:batch_predict(rows,target)
    temporary.replace(args.output)
    evidence={"documents":documents,"windows":windows,"seconds":time.perf_counter()-started,"backend":"ONNX Runtime CPU","batch_size":args.batch_size,"raw_unfiltered":True}
    args.output.with_suffix(".evidence.json").write_text(json.dumps(evidence,indent=2)+"\n")
    print(json.dumps(evidence),flush=True)

if __name__=="__main__":main()