"""Recognize one cropped word or line with an explicit language and text style."""
import argparse
import json
from pathlib import Path
import sys


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('image', type=Path)
    parser.add_argument('--language', choices=['ne', 'hi', 'en'], required=True)
    parser.add_argument('--text-style', choices=['handwriting', 'print'], default='handwriting',
                        help='English print uses the character branch; default: handwriting.')
    parser.add_argument('--backend', choices=['torch', 'paddle'], default='torch')
    parser.add_argument('--model-dir', type=Path, default=Path(__file__).resolve().parent)
    args = parser.parse_args()
    root = args.model_dir.resolve()
    branch = 'english' if args.language == 'en' and args.text_style == 'handwriting' else 'character'
    if args.backend == 'torch':
        import torch
        from PIL import Image
        from transformers import TrOCRProcessor, ViTImageProcessor
        sys.path.insert(0, str(root / 'code'))
        from shared_character_recognizer import SharedCharacterRecognizer
        from trocr_character_codec import CharacterCodec
        torch.set_num_threads(2)
        torch.use_deterministic_algorithms(True)
        folder = root / 'source-model'
        model = SharedCharacterRecognizer.from_bundle(folder)
        assets = folder / ('english-assets' if branch == 'english' else 'character-assets')
        processor = (TrOCRProcessor.from_pretrained(assets, local_files_only=True, use_fast=False)
                     if branch == 'english' else ViTImageProcessor.from_pretrained(assets, local_files_only=True))
        with Image.open(args.image) as image:
            pixels = processor(images=image.convert('RGB'), return_tensors='pt').pixel_values
        generation = dict(do_sample=False, num_beams=10, max_new_tokens=256 if branch == 'english' else 257,
                          min_length=0, length_penalty=1., early_stopping=False, use_cache=True)
        if branch == 'character':
            generation['suppress_tokens'] = [0, 1, 3]
        with torch.inference_mode():
            tokens = model.generate_pixels(pixels, branch, generation)[0].tolist()
        if branch == 'english':
            text = processor.batch_decode([tokens], skip_special_tokens=True)[0]
        else:
            spec = json.loads((assets / 'character_codec.json').read_text())
            text = CharacterCodec(list(spec['character_ids'])).decode(tokens)
        result = dict(text=text, tokens=tokens, token_limit_reached=len(tokens)-1 >= generation['max_new_tokens'] and tokens[-1] != 2)
    else:
        import paddle
        from PIL import Image
        sys.path.insert(0, str(root / 'paddle/code'))
        from shared_character_native import SharedCharacterNative, preprocess_rgb
        paddle.set_device('cpu')
        folder = root / 'paddle'
        model = SharedCharacterNative(folder / 'candidate', folder / 'encoder', folder / 'decoders',
                                     allow_unqualified=True, library_path=folder / 'encoder-library.so')
        with Image.open(args.image) as image:
            result = model.predict_pixels(preprocess_rgb(image), branch)
    print(json.dumps(dict(language=args.language, text_style=args.text_style, branch=branch, backend=args.backend, **result), ensure_ascii=False))


if __name__ == '__main__':
    main()