"""Standalone quantized line recognition; the page pipeline is a separate layer."""
import argparse,collections,hashlib,json
from pathlib import Path
import numpy as np
from rtl import flip

class Recognizer:
    def __init__(self,root=None,role='body',provider='CPUExecutionProvider',cache_entries=512,variant='block16'):
        self.root=Path(root or Path(__file__).resolve().parent);self.role=role
        config=json.loads((self.root/'quantization_config.json').read_text());self.backend=config['format'];self.variant=config.get('default_variant','block16')
        self.alphabet=json.loads((self.root/'vocab.json').read_text())['chars'];d=self.root/'components'/role
        self.cache=collections.OrderedDict();self.cache_entries=cache_entries;self.calls=0;self.cache_hits=0
        if self.backend=='onnx':
            import onnxruntime as ort
            if provider not in ort.get_available_providers():raise RuntimeError('Unavailable provider: '+provider)
            options=ort.SessionOptions();options.intra_op_num_threads=2;options.inter_op_num_threads=1
            if variant=='hqq32'and provider!='CPUExecutionProvider':raise ValueError('HQQ32 helper is qualified for ORT CPU only')
            name='recognizer-hqq4-block32.onnx'if variant=='hqq32'else'recognizer-w4-block16.onnx'
            self.session=ort.InferenceSession(str(d/'models'/name),sess_options=options,providers=[provider]);self.run=lambda x,w:self.session.run(None,dict(pixels=x,widths=w))
        elif self.backend=='coreml':
            import coremltools as ct
            self.session=ct.models.MLModel(str(d/'models/recognizer-w4-block16.mlpackage'),compute_units=ct.ComputeUnit.CPU_AND_GPU)
            def run(x,w):
                o=self.session.predict(dict(pixels=x,widths=w.astype(np.int32)));return o['tokens'],o['confidence'],o['lengths']
            self.run=run
        elif self.backend=='litert':
            from ai_edge_litert.interpreter import Interpreter
            self.session=Interpreter(model_path=str(d/'litert/recognizer-dynamic4-block32.tflite'),num_threads=2);self.runners={}
            def run(x,w):
                width=x.shape[-1]
                if width not in self.runners:self.runners[width]=self.session.get_signature_runner('width'+str(width))
                o=self.runners[width](args_0=x,args_1=w.astype(np.int32));return o['output_0'],o['output_1'],o['output_2']
            self.run=run
        else:raise ValueError(self.backend)
    def predict_ink(self,ink):
        ink=np.asarray(ink,dtype=np.float32)
        if ink.ndim!=2 or ink.shape[0]!=64 or ink.shape[1]<1 or not np.isfinite(ink).all():raise ValueError('Finite float32 ink map [64,true_width] required')
        key=hashlib.sha256(str(ink.shape).encode()+ink.tobytes()).hexdigest()
        if key in self.cache:
            self.cache_hits+=1;self.cache.move_to_end(key);return dict(self.cache[key])
        width=ink.shape[1];bucket=max(512,2**int(np.ceil(np.log2(width))))
        if bucket>8192:raise ValueError('Width exceeds the validated8192 bound; never truncate a line')
        x=np.zeros((1,1,64,bucket),np.float32);x[0,0,:,:width]=ink
        tokens,confidence,lengths=self.run(x,np.array([width],np.int64));self.calls+=1
        n=int(lengths[0]);seq=np.asarray(tokens[0,:n]);keep=(seq!=0)&np.r_[True,seq[1:]!=seq[:-1]];ids=seq[keep].tolist()
        visual=''.join(self.alphabet[int(i)-1]for i in ids)
        result=dict(text=flip(visual,base_rtl=True),visual_text=visual,confidence=float(np.asarray(confidence[0,:n])[keep].mean())if keep.any()else 0.,role=self.role)
        if self.cache_entries>0:
            self.cache[key]=result
            while len(self.cache)>self.cache_entries:self.cache.popitem(last=False)
        return dict(result)
    def read_line(self,image,box=None):
        from geometry import crop_line64
        gray=np.asarray(image.convert('L'));box=box or [0,0,gray.shape[1],gray.shape[0]]
        ink,receipt=crop_line64(gray,box);result=self.predict_ink(ink);result['geometry']=receipt;return result

if __name__=='__main__':
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('line_image');p.add_argument('--role',default='body',choices=['body','short_labels','printed','numeric_math','degraded']);p.add_argument('--provider',default='CPUExecutionProvider');p.add_argument('--variant',default='block16',choices=['block16','hqq32']);a=p.parse_args()
    from PIL import Image
    reader=Recognizer(role=a.role,provider=a.provider,variant=a.variant);print(json.dumps(reader.read_line(Image.open(a.line_image)),ensure_ascii=False,indent=2))