"""Process-local PaddleX adapter for the shared two-branch research recognizer."""
import math
from pathlib import Path

import numpy as np
from PIL import Image
from paddlex.inference.models.predictors.local_model_predictor import LocalModelPredictor
from paddlex.inference.models.text_recognition.predictor import ImageBatchSampler,TextRecResult
from paddlex.inference.models.bindings.registry import register_predictor_binding_map

from shared_character_native import SharedCharacterNative,preprocess_rgb,sha


MODEL_NAME='TrOCR_shared_character_Paddle_research_rec'


class PaddleXSharedCharacter(LocalModelPredictor):
    def __init__(self,*args,research_only_allow_unqualified=False,**kwargs):
        super().__init__(*args,**kwargs)
        if self.model_name!=MODEL_NAME or self.device!='cpu':raise ValueError('Require the shared CPU research model')
        native=self.config['Native']
        if native['format']!='shared-character-paddle-research-v1' or native['parameters']!=74842752:
            raise ValueError('Unexpected shared recognizer profile')
        for name,digest in native['files_sha256'].items():
            relative=Path(name)
            if relative.is_absolute() or '..' in relative.parts or relative.as_posix()!=name:raise ValueError('Invalid bundle asset path')
            if sha(self.model_dir/relative)!=digest:raise ValueError('Shared bundle asset bytes changed')
        self.runtime=SharedCharacterNative(self.model_dir/'candidate',self.model_dir/'encoder',self.model_dir/'decoders',
            allow_unqualified=research_only_allow_unqualified,library_path=self.model_dir/'encoder-library.so')

    def _build_batch_sampler(self):return ImageBatchSampler()
    def _get_result_class(self):return TextRecResult

    def process(self,batch_data,branch=None,return_word_box=False):
        if branch not in ['english','character']:raise ValueError('Explicit research branch is required')
        if return_word_box:raise ValueError('Autoregressive word-box alignment is not implemented')
        images=[];outputs=[]
        for item in batch_data.instances:
            if isinstance(item,(str,Path)):
                with Image.open(item) as source:image=source.convert('RGB')
                bgr=np.asarray(image)[:,:,::-1].copy()
            elif isinstance(item,np.ndarray) and item.dtype==np.uint8 and item.ndim==3 and item.shape[2]==3:
                bgr=item;image=Image.fromarray(np.ascontiguousarray(bgr[:,:,::-1]))
            else:raise ValueError('Require an image path or uint8 HWC BGR array')
            images.append(bgr);outputs.append(self.runtime.predict_pixels(preprocess_rgb(image),branch))
        return dict(input_path=batch_data.input_paths,page_index=batch_data.page_indexes,input_img=images,
            rec_text=[r['text'] for r in outputs],rec_score=[min(1.,math.exp(r['score'])) for r in outputs],
            beam_score=[r['score'] for r in outputs],tokens=[r['tokens'] for r in outputs],
            decoder_iterations=[r['decoder_iterations'] for r in outputs],
            token_limit_reached=[r['token_limit_reached'] for r in outputs],recognizer_branch=[branch]*len(outputs),
            score_kind=['Uncalibrated geometric mean of generated token probabilities']*len(outputs),vis_font=[None]*len(outputs))


def register():register_predictor_binding_map(PaddleXSharedCharacter,{'flexible':[MODEL_NAME]})