import argparse
import sys
from pathlib import Path
import unicodedata
import torch
from PIL import Image, ImageOps
from model import CRNN


def predict(image_path, checkpoint_path):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True)
    vocab = checkpoint['vocab']
    reverse = {i: c for c, i in vocab.items()}
    if '<UNK>' in vocab:
        reverse[vocab['<UNK>']] = '\ufffd'
    model = CRNN(len(vocab))
    model.load_state_dict(checkpoint['model_state_dict'], strict=True)
    model.to(device).eval()
    height = checkpoint['height']
    with Image.open(image_path) as source:
        image = source.convert('L')
    width = max(4, round(image.width * height / image.height))
    image = ImageOps.mirror(image.resize((width, height), Image.Resampling.LANCZOS))
    data = torch.frombuffer(bytearray(image.tobytes()), dtype=torch.uint8)
    data = (data.reshape(1, 1, height, width).float() / 255.0 - 0.5) / 0.5
    with torch.inference_mode():
        ids = model(data.to(device)).argmax(-1)[0].cpu().tolist()
    result, previous = [], None
    for index in ids:
        if index != vocab['<BLANK>'] and index != previous:
            result.append(reverse[index])
        previous = index
    return unicodedata.normalize('NFC', ''.join(result))


if __name__ == '__main__':
    if hasattr(sys.stdout, 'reconfigure'):
        sys.stdout.reconfigure(encoding='utf-8')
    parser = argparse.ArgumentParser()
    parser.add_argument('image')
    parser.add_argument('--checkpoint', default=str(Path(__file__).parent/'best_cer_v10.pt'))
    args = parser.parse_args()
    print(predict(args.image, args.checkpoint))