import cv2
import sys
from PIL import Image
import matplotlib.pyplot as plt
from pytesseract import image_to_string
from g_translate import translate
import codecs

# First check if cmd argument includes at least target image file path.
if len(sys.argv) < 2:
    print("Usage: Too few arguments! You need to specify at least the target image file.")
    sys.exit(0)

# Prefix to be used for all output files generated by this program.
IMG_PREFIX = sys.argv[1][:sys.argv[1].rindex(".")]

# Languages to be translated into. May be empty.
DEST = [sys.argv[i] for i in range(2, len(sys.argv))]

# Johns Hopkins Logo to use in template matching.
JHU_LOGO = cv2.imread('data/jhu_logo.png', 0)

# dict: {Google Translate API language code : English expression for that language}.
# Used for stdout pretty print.
LANGUAGES = {"en": "English", "fr": "French", "es": "Spanish", "ko": "Korean", \
             "ja": "Japanese", "zh-CN": "Chinese (traditional)", "ar": "Arabic", \
             "de": "German", "hi": "Hindi", "ru": "Russian", "el": "Greek"}

"""
Greets users of this program.
"""
def greetings():
    print("======= Welcome to James' text recognition and translation software! =======")

"""
Thresholds the input image file and outputs the result as an image file.
Blurring and denoising applied to image to provide clean input for OCR.
"""
def threshold_image():
    img = cv2.imread(sys.argv[1], 0)
    img = cv2.medianBlur(img, 5)
    th_img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\
                                   cv2.THRESH_BINARY, 5, 2)
    th_img = cv2.fastNlMeansDenoising(th_img, None, 7, 21, 7)
    cv2.imwrite(IMG_PREFIX + '_thresh.png', th_img)

"""
Call pytesseract's ocr function (image_to_string) and process output.

return:
    raw_text - str: direct output of pytesseract image_to_string.
    split_text - list[str]: list of words that appear in raw_text.
"""
def extract_text():
    raw_text = image_to_string(Image.open(IMG_PREFIX + '_thresh.png'))
    return raw_text

"""
Write output text file for input image.
"""
def write_original_output(raw_text):
    split_text = raw_text.split(" ")
    print("Total of " + str(len(split_text)) + " words detected.")
    with open(IMG_PREFIX + '_output.txt', 'w') as writer:
        try:
            for word in split_text:
                writer.write(word + " ")
        except IOError:
            raise Exception("Exception while writing text file for input image "\
                            + sys.argv[1] + ".")    
    print("Output successfully written in original language.")

"""
Translate OCR recognized text into target languages supplied by cmd line args.
String raw_text chopped by 1000 chars to prevent Google Translate API 400 Error.
Write output text file for each target language in utf-8.
"""
def write_translated_output(text, src):
    # Chop text by 1000 characters each.
    chopped_text = []
    for i in range(0, len(text), 1000):
        end_index = i + 1000 if i + 1000 < len(text) else len(text)
        chopped_text.append(text[i:end_index])

    # Translate into each target language.
    for dest_lang in DEST:
        for c in chopped_text:
            translated = translate(c, dest_lang, src).decode('utf-8')
            try:
                codecs.open(IMG_PREFIX + '_output_' + dest_lang \
                            + '.txt', 'a', 'utf-8').write(translated)
            except IOError:
                raise Exception("Exception while writing text file for language "\
                                + LANGUAGES[dest_lang] + ".")
        print("Output successfully written in " + LANGUAGES[dest_lang] + ".")

"""
Code adapted from OpenCV Template Matching tutorial.

Takes the JHU logo in 5 different scales and tries to match this logo with anything
that appears in the input document image. If for any of the 5 templates the matchTemplate
function returns a value greater than 0.51 (chosen empirically), we consider this as
a positive match and let the users know that this is a JHU document.
"""
def check_jhu_document(do_plot=False):
    img = cv2.imread(sys.argv[1], 0)
    is_jhu = False
    # JHU logo templates.
    logos = ['data/jhu_logo1.png', 'data/jhu_logo2.png', 'data/jhu_logo3.png', \
             'data/jhu_logo4.png', 'data/jhu_logo5.png',]
    for template in logos:
        img2 = img.copy()
        JHU_LOGO = cv2.imread(template, 0)
        method = eval('cv2.TM_CCOEFF_NORMED')
        try:
            res = cv2.matchTemplate(img2, JHU_LOGO, method)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
            if max_val > 0.51:
                is_jhu = True
            if do_plot:
                top_left = max_loc
                w, h = JHU_LOGO.shape[::-1]
                bottom_right = (top_left[0] + w, top_left[1] + h)
                cv2.rectangle(img2, top_left, bottom_right, (0,0,255), 2)
                plt.subplot(121),plt.imshow(res,cmap='gray')
                plt.title('Matching Result'), plt.xticks([]), plt.yticks([])
                plt.subplot(122),plt.imshow(img2,cmap='gray')
                plt.title('Detected Point'), plt.xticks([]), plt.yticks([])
                plt.show()
        except:
            pass
    if is_jhu:
        print("This is a JHU document!!! Go Hopkins!")
    else:
        print("Unfortunately, I don't think there's enough evidence to conclude that this")
        print("is a JHU docoument...")

"""
Bids users farewell.
"""
def farewell():
    print("================================ Good Bye! =================================")
    return 0

"""
Main driver function for this program.
"""
def main():
    greetings()
    threshold_image()
    raw_text = extract_text()    
    write_original_output(raw_text)
    write_translated_output(raw_text, "auto")
    check_jhu_document()
    return farewell()

if __name__ == '__main__':
    main()
