from skimage.io import imread, imsave
from skimage.color import rgb2gray
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from skimage.transform import resize
from utils.preprocessing import preprocessImage, postProcessImage, process_segment_and_crop_image
from utils.line_segmentation import segment_image_to_lines
from configs import unet_enabled
from utils.helper import load_images_from_json
from inference import Inference
from configs import model_path, processor_path, unet_model_path
from utils.postprocessing import PostProcessing
from loguru import logger

def RenAITranscription(image, post_processing_enabled=False,unet_enabled=False):
   # 1- preprocessing
    org_img = imread(image)[: , : ,:]

    logger.info(f'Image Dimensions : {org_img.shape[0]} x {org_img.shape[1]}')

    intial_process_image = preprocessImage(org_img)

    if unet_enabled:
        logger.info("Masked based segmentation and cropping enabled...")    
        cropped_img = process_segment_and_crop_image(unet_model_path, org_img, intial_process_image, padding=10, min_contour_area=100)
        processed_image = postProcessImage(cropped_img)
        logger.info(f"Image cropped and Pre-processed successfully.....")
    else:
        logger.info("Image Preprocessing started......")
        processed_image = postProcessImage(intial_process_image)
        logger.info(f"Image Pre-processed successfully.....")
    
    # 2 - Line segmentation Algorithm
    line_segments = segment_image_to_lines(processed_image, base_key="line",ct=0)

    # 3 - Model Inference

    transciption_generator = Inference(
            model_path=model_path,
            processor_path=processor_path,
            target_size=(256, 64),
            batch_size=32
    )
    result = transciption_generator.generate_texts_from_images(line_segments)

    # Generated texts
    for key, value in result.items():
        print(f"{key}: {value['transcription']}")

    # 4 - Post processing
    # Dictionary based fuzzy matching
    if post_processing_enabled:
        for key, value in result.items():
            corrected = PostProcessing(value['transcription'])
            result[key]['post_processed'] = corrected
            print(f"{key}: {value['post_processed']}")

    print(result)

    logger.info("Transcription completed successfully!")    
    return result

if __name__ == "__main__":
    RenAITranscription("1.png", post_processing_enabled=False, unet_enabled=False)