from PIL import Image
import numpy as np
import math
import numpy as np
import onnxruntime as ort
import os

def resize_norm_img(img):
    imgC, imgH, imgW = 3, 48, 320
    h, w = img.shape[:2]
    ratio = w / float(h)
    resized_w = imgW if math.ceil(imgH * ratio) > imgW else int(math.ceil(imgH * ratio))

    pil_img = Image.fromarray(img)
    resized_image = np.array(pil_img.resize((resized_w, imgH), Image.BILINEAR)).astype('float32')

    resized_image = resized_image.transpose((2, 0, 1)) / 255.0
    resized_image -= 0.5
    resized_image /= 0.5

    padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
    padding_im[:, :, 0:resized_w] = resized_image
    return np.expand_dims(padding_im, axis=0)

def load_dict(dict_path):
    """Loads the dictionary and aligns it with PaddleOCR's CTC index format."""
    character = []
    with open(dict_path, "rb") as fin:
        lines = fin.readlines()
        for line in lines:
            line = line.decode('utf-8').strip("\n").strip("\r\n")
            character.append(line)
            
    character.append(" ")
    
    # Reserve index 0 for the blank token
    character = ['<blank>'] + character
    return character

def decode(preds, character):
    """Converts the raw ONNX matrix output into readable text using CTC rules."""
    # Extract the highest probability class indices
    preds_idx = preds.argmax(axis=2)[0]
    preds_prob = preds.max(axis=2)[0]
    
    char_list = []
    conf_list = []
    
    for i in range(len(preds_idx)):
        # CTC rules: Ignore the blank token (0) and ignore consecutive duplicate characters
        if preds_idx[i] != 0 and not (i > 0 and preds_idx[i] == preds_idx[i - 1]):
            char_list.append(character[preds_idx[i]])
            conf_list.append(preds_prob[i])
            
    text = ''.join(char_list)
    confidence = np.mean(conf_list) if len(conf_list) > 0 else 0.0
    return text, confidence

# Example usage
if __name__ == "__main__":
    session = ort.InferenceSession(
        "model.onnx", 
        providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
    )
    
    input_name = session.get_inputs()[0].name
    
    # 2. Load Greek dictionary
    char_dict = load_dict("greek_dict.txt")
    
    # 3. Read and preprocess the test image
    for file in os.listdir("test_images"):
        if file.endswith(".jpeg") or file.endswith(".jpg") or file.endswith(".png"):
            img = np.array(Image.open(os.path.join("test_images", file)).convert("RGB"))
            img = img[:, :, ::-1]  # RGB → BGR to match PaddleOCR training data
            input_tensor = resize_norm_img(img)

            outputs = session.run(None, {input_name: input_tensor})
            raw_predictions = outputs[0]
            
            text, confidence = decode(raw_predictions, char_dict)
            
            print(f"Recognized Text from {file}: {text}")
            print(f"Confidence Score: {confidence:.4f}")