# -*- coding: utf-8 -*-
r"""

Ce script contient des fonctions pour océriser (numériser) une image, effectuer des corrections automatiques,
identifier les noms des personnages et les didascalies, et marquer les blocs reconnus par l'outil d'OCR dans l'image,
avec le niveau de confiance de la reconnaissance.

Informations utiles :
- https://nanonets.com/blog/ocr-with-tesseract/
- http://info.blaisepascal.fr/tesseract

1. Installer Tesseract à partir de : 
        https://github.com/UB-Mannheim/tesseract/wiki
   Lors du processus d'installation, installer également les composants 
   suivants (cocher les cases) :
       Additional language data : French et German
   Sous Windows, Tesseract est normalement installé dans le dossier:
       C:\Users\nom_utilisateur\AppData\Local\Tesseract-OCR
2. Créer un nouvel environnement appelé 'ocr' pour python 3.7 avec les 
librairies suivantes :
   - pytesseract
   - pillow
   - opencv
   - numpy
    
"""
from PIL import Image
import re
import argparse
import cv2
import pytesseract
import numpy as np
from pytesseract import Output
import sys
TESSERACT_DIR = r'C:/Program Files/Tesseract-OCR/tesseract.exe'
CUSTOM_CONFIG = r'-l deu+fra --psm 6'
BORDER_COLOR_GREEN = (0, 255, 0)
BORDER_COLOR_RED = (0, 0, 255)
img = Image.open('059.jpg')
d = pytesseract.image_to_data(img, output_type=Output.DICT)
txt = pytesseract.image_to_string(img, config=CUSTOM_CONFIG)

def get_dimensions(i, image_data):
    """
    Cette fonction prend pour argument un entier désignant un indice de position et un dictionnaire de données picturales
    et retourne les dimensions du cadre englobant selon l'indice sélectionné (marge gauche, marge haute, largeur, hauteur) obtenues avec
    """
    left = image_data['left'][i]
    top = image_data['top'][i]
    width = image_data['width'][i]
    height = image_data['height'][i]
    return left, top, width, height


def frame_confidence(img, image_data):
    """Cette fonction prend pour argument un fichier image et le dictionnaire contenant ses données (les mots, les
    dimensions de leur cadre, le texte contenu, le niveau de confiance de la reconnaissance) et retourne un encadrement
    en vert pour les mots dont le niveau de confiance est supérieur à 80 et rouge pour le reste"""
    new_img = np.copy(img)
    n_boxes = len(image_data['text'])
    for i in range(n_boxes):
        (x, y, w, h) = get_dimensions(i, image_data)
        if int(image_data['conf'][i]) > 80:
            new_img = cv2.rectangle(new_img, (x, y), (x + w, y + h), 
                                    BORDER_COLOR_GREEN, 1)
        else:
            new_img = cv2.rectangle(new_img, (x, y), (x + w, y + h), 
                                    BORDER_COLOR_RED, 1)
    return new_img

def get_text(img):
    """retourne le contenu textuel du fichier image"""
    return pytesseract.image_to_string(img, config=CUSTOM_CONFIG)

def annotate_speaker_names(text):
    """Cette fonction prend pour argument le contenu texte d'une image et renvoie le texte avec
    les noms de personnages entourés d'étoiles"""
    regex_2 = re.compile(r"^.*:")
    regex_3 = re.compile(r"^\w+:$")
    regex_4 = re.compile(r"^\w+\):$")
    regex_5 = re.compile(r"^\(.*\):$")
    split = text.split("\n")
    for elem in range(len(split)):
        if regex_2.search(split[elem]):
            split_2 = split[elem].split(" ")
            for element in split_2:
                if regex_3.search(split_2[0]):
                    element = " ".join(split_2).replace(split_2[0], "**"+split_2[0][:-1]+"**"+":")
                    split[elem] = element
                elif regex_3.search(split_2[1]):
                    element = " ".join(split_2).replace(split_2[1],split_2[1][:-1]+"**"+":").replace(split_2[0],"**"+split_2[0])
                    split[elem] = element
                elif regex_4.search(element):
                    element = " ".join(split_2).replace(element, element[:-1]+":").replace(split_2[0], "**"+split_2[0]+"**")
                    split[elem] = element
                elif regex_5.search(element):
                    element = " ".join(split_2).replace(element, "**"+element[:-1]+":").replace(split_2[0],"**"+split_2[0])
                    split[elem] = element
    nv_split = "\n".join(split)
    annotated_text = nv_split
    return annotated_text

def annotate_stage_directions(text):
    """Cette fonction prend en argument du contenu texte et le renvoie avec les didascalies entourées de tirets-bas"""
    annotated_text = text.replace("(", "_(").replace(")", ")_")
    return annotated_text


def read_corrections_file(corrections_file):
    """Cette fonction lit le fichier de corrections et renvoie une liste composée de listes comprenant la forme erronnée
    et de la forme correcte"""
    corrections = []
    with open(corrections_file, 'r', encoding='utf8') as f:
        for line in f:
            corrections.append(list(line.rstrip().split("\t")))
    return corrections

def correct_text(text, corrections_file):
    """Cette fonction prend pour argument le contenu texte de l'image et le fichier de corrections et renvoie
    le texte corrigé"""
    with open(corrections_file, 'r', encoding='utf8') as f:
        for line in f:
            line = line.rstrip().split("\t")
            text = text.replace(line[0], line[1])
    corrected_text = text
    return corrected_text

def get_data(img):
    """Cette fonction renvoie les données verbeuses de l'image (boîtes, les degrés de confiance, les numéros de
    ligne et de page) sous forme de dictionnaire"""
    return pytesseract.image_to_data(img, output_type=pytesseract.Output.DICT, 
                                     config=CUSTOM_CONFIG)

def save_image(img, output_file):
    """Cette fonction sauvegarde le fichier image modifié sous un autre nom/extension"""
    cv2.imwrite(output_file, img)
    
def save_text(text, output_file):
    """Cette fonction sauvegarde le contenu texte du fichier image dans un nouveau fichier texte"""
    with open(output_file, 'w', encoding='utf8') as f:
        f.write(text)

def demo(img_path, corrections_file, output_img_file, output_text_file):
    """Cette fonction océrise une image, effectue des corrections automatiques, identifie les noms des personnagges et
    les didascalies et marque les blocs reconnus dans l'image avec le niveau de confiance de la reconnaissance"""
    img = cv2.imread(img_path)
    
    d = get_data(img)
    new_img = frame_confidence(img, d)
    save_image(new_img, output_img_file)
    
    text = get_text(img)
    text = correct_text(text, corrections_file)
    text = annotate_speaker_names(text)
    text = annotate_stage_directions(text)
    save_text(text, output_text_file)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("image", help="path to image")
    parser.add_argument("corrections", help="path to corrections file")
    parser.add_argument("output_image", help="path to output image")
    parser.add_argument("output_text", help="path to output text")
    args = parser.parse_args()
    pytesseract.pytesseract.tesseract_cmd = TESSERACT_DIR
    demo(args.image, args.corrections, args.output_image, args.output_text)

if __name__ == "__main__":
    main()