#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright © Zdenko Podobný 2014-2017
# Licensed under the terms of the Apache License Version 2.0

"""
   Simple python demo of tesseract-ocr
"""

# https://www.mail-archive.com/pyqt@riverbankcomputing.com/msg20541.html
# http://pastebin.com/DhPUgrAj

__author__ = 'Zdenko Podobný <zdenop@gmail.com>'
__version__ = '0.3'
__date__ = '21.05.2017'

import os
import sys
import ctypes
import locale

from PyQt5.QtWidgets import (QApplication, QMainWindow, QStyleFactory,
                             QFileDialog, QStyle, QGraphicsScene)
from PyQt5.QtGui import (QTextCursor, QBrush, QColor, QTransform, QPixmap,
                         QPen, QIcon)
from PyQt5.QtCore import (Qt, pyqtSignal, pyqtSlot, QCoreApplication, QEvent,
                          QRectF)

from ui.ui_mainwindow import Ui_MainWindow

import libs.tesstool as tess
import libs.lepttool as lept
import libs.settings as sett

from libs.scene import CustomGraphicsScene
from libs.areaitem import AreaItem


class MainWindow(QMainWindow, Ui_MainWindow):
    """Class For MainWindow
    """
    dropped_to_scene = pyqtSignal('QString')

    def __init__(self, parent=None):
        """ Constructor
        """
        super(MainWindow, self).__init__(parent)
        self.setAttribute(Qt.WA_DeleteOnClose)
        QApplication.setStyle(QStyleFactory.create('cleanlooks'))
        self.setupUi(self)
        self.scene = CustomGraphicsScene()
        self.scene.dropped.connect(self.load_image)
        self.graphicsView.setScene(self.scene)
        self.scene.installEventFilter(self)
        self.graphicsView.setBackgroundBrush(QBrush(Qt.gray, Qt.BDiagPattern))
        quit_icon = QApplication.style().standardIcon(
            QStyle.SP_DialogCloseButton)
        self.pushButtonQuit.setIcon(quit_icon)
        self.setWindowTitle('Analyze & ORC image with tesseract and leptonica')
        self.actionZoomOut.triggered.connect(self.zoomOut)
        self.actionZoomIn.triggered.connect(self.zoomIn)
        self.actionZoomTo1.triggered.connect(self.zoomTo1)
        self.actionZoomFit.triggered.connect(self.zoomFit)
        self.actionZoomToWidth.triggered.connect(self.zoomToWidth)

        # Initialize variables and pointers
        self.box_data = []
        self.pix_image = False
        self.image = None
        self.image_width = 0
        self.image_height = 0
        self.tesseract = None
        self.api = None
        self.lang = 'eng'

        lang = sett.readSetting('language')
        if lang:
            self.lang = lang

        self.initialize_tesseract()
        if self.tesseract:
            available_languages = tess.get_list_of_langs(self.tesseract,
                                                         self.api)
            if available_languages:
                for lang in available_languages:
                    self.comboBoxLang.addItem(lang, lang)
            current_index = self.comboBoxLang.findData(self.lang)
            if current_index:
                self.comboBoxLang.setCurrentIndex(current_index)

        for idx, psm in enumerate(tess.PSM):
            self.comboBoxPSM.addItem(psm, idx)
        for idx, ril in enumerate(tess.RIL):
            self.comboBoxRIL.addItem(ril, idx)

        self.leptonica = lept.get_leptonica()
        if not self.leptonica:
            self.show_msg('Leptonica initialization failed...')

        # Read settings and set default values
        geometry = sett.readSetting('settings_geometry')
        if geometry is not None:
            try:
                self.restoreGeometry(geometry)
            except TypeError:
                pass
        else:
            self.resize(1150, 950)
        state = sett.readSetting('state')
        if state is not None:
            self.restoreState(state)
        sp_1_state = sett.readSetting('splitter_1Sizes')
        if sp_1_state is not None:
            try:
                self.splitter_1.restoreState(sp_1_state)
            except TypeError:
                pass
        sp_2_state = sett.readSetting('splitter_2Sizes')
        if sp_2_state is not None:
            try:
                self.splitter_2.restoreState(sp_2_state)
            except TypeError:
                pass
        psm = sett.readSetting('PSM')
        if psm:
            current_index = self.comboBoxPSM.findData(psm)
            self.comboBoxPSM.setCurrentIndex(current_index)
        ril = sett.readSetting('RIL')
        if ril:
            current_index = self.comboBoxRIL.findData(ril)
            self.comboBoxRIL.setCurrentIndex(current_index)

        image_name = sett.readSetting('images/last_filename')
        if image_name:
            self.image_name = image_name
        self.load_image(image_name)
        zoom_factor = sett.readSetting('images/zoom_factor')
        self.setZoom(zoom_factor)

    def initialize_tesseract(self):
        """Create tesseract api
        """
        self.tesseract = tess.get_tesseract(os.path.dirname(__file__))
        if not self.tesseract:
            self.show_msg('Tesseract initialization failed...')
            return
        self.api = self.tesseract.TessBaseAPICreate()
        tessdata_prefix = tess.get_tessdata_prefix()
        # current_locale = locale.getlocale()  # Save current locale
        # Switch to C locale to handle
        #  Error: Illegal min or max specification!
        #  "Fatal error encountered!" == NULL:Error:Assert failed:in file
        #   ../../tesseract-ocr/ccutil/globaloc.cpp, line 75
        locale.setlocale(locale.LC_ALL, 'C')
        retc = self.tesseract.TessBaseAPIInit3(self.api,
                                               tessdata_prefix.encode(),
                                               self.lang.encode())
        # Restore saved locale
        # locale.setlocale(locale.LC_ALL, current_locale)
        if (retc):
            self.tesseract.TessBaseAPIDelete(self.api)
            self.show_msg('<span style="color:red">Could not initialize '
                          'tesseract.</span>')
            return
        self.show_msg('Tesseract {0} initialized with language \'{1}\'.\n'
                      '{2} is used with image support for {3}'
                      .format(tess.VERSION, self.lang,
                              lept.get_version(), lept.get_image_support()))

    @pyqtSlot()
    def on_pushButtonShow_pressed(self):
        """Display rectangles
        """
        tessdata_prefix = tess.get_tessdata_prefix()
        locale.setlocale(locale.LC_ALL, 'C')
        retc = 0
        if self.lang != str(self.comboBoxLang.currentText()):
            self.lang = str(self.comboBoxLang.currentText())
            retc = self.tesseract.TessBaseAPIInit3(self.api,
                                                   tessdata_prefix,
                                                   self.lang)
            self.show_msg('Using language \'%s\'.' % self.lang)
        if (retc):
            self.tesseract.TessBaseAPIDelete(self.api)
            self.show_msg('<span style="color:red">Could not re-initialize '
                          'tesseract.</span>')
            return

        # Shut up tesseract - there could be a lot of unwanted messages
        #tesseract.TessBaseAPISetVariable(api, "debug_file", "/dev/null")

        self.tesseract.TessBaseAPIClear(self.api)
        # Set PIX structure to tesseract api
        self.tesseract.TessBaseAPISetImage2(self.api, self.pix_image)

        self.tesseract.TessBaseAPISetPageSegMode(
            self.api, self.comboBoxPSM.currentIndex())

        # Get info(BOXA structure) about lines(RIL_TEXTLINE) from image in api
        boxa = self.tesseract.TessBaseAPIGetComponentImages(
            self.api, self.comboBoxRIL.currentIndex(), 1, None, None)
        if not boxa:
            self.show_msg('No component found. Try to change PSM or RIL.')
            return

        # Get info about number of items on image
        n_items = self.leptonica.boxaGetCount(boxa)
        psm = tess.PSM[self.comboBoxPSM.currentIndex()]
        ril = tess.RIL[self.comboBoxRIL.currentIndex()]
        self.show_msg('<span style="color:green">'
                      'Found %d image components with %s and %s.</span>' %
                      (n_items, psm, ril))

        ocr_psm = tess.PSM_SINGLE_BLOCK
        if ril == 'RIL_PARA':
            ocr_psm = tess.PSM_SINGLE_BLOCK
        elif ril == 'RIL_TEXTLINE':
            ocr_psm = tess.PSM_SINGLE_LINE
        elif ril == 'RIL_WORD':
            ocr_psm = tess.PSM_SINGLE_WORD
        elif ril == 'RIL_SYMBOL':
            ocr_psm = tess.PSM_SINGLE_CHAR
        self.tesseract.TessBaseAPISetPageSegMode(self.api, ocr_psm)

        # Set up result type (BOX structure) for leptonica function boxaGetBox
        self.leptonica.boxaGetBox.restype = lept.BOX_PTR_T
        self.leptonica.boxaGetBox.argtypes = []
        n_boxes = len(self.box_data)
        if n_boxes:
            for idx in range(n_boxes):
                self.scene.removeItem(self.box_data[idx])

        # Display items and print its info
        box_items = []
        for item in range(0, n_items):
            lept_box = self.leptonica.boxaGetBox(boxa, item, lept.L_CLONE)
            box = lept_box.contents
            self.tesseract.TessBaseAPISetRectangle(self.api,
                                                   box.x, box.y, box.w, box.h)
            ocr_result = self.tesseract.TessBaseAPIGetUTF8Text(self.api)
            result_text = ctypes.string_at(ocr_result).decode('utf-8').strip()
            conf = self.tesseract.TessBaseAPIMeanTextConf(self.api)
            self.show_msg('Box[%d]: x=%d, y=%d, w=%d, h=%d, '
                          'confidence: %d, '
                          'text: <span style="color:blue">%s</span>' %
                          (item, box.x, box.y, box.w, box.h,
                           conf, result_text
                           ))
            box_items.append(
                self.scene.addRect(
                    box.x, box.y, box.w, box.h, QPen(
                        QColor(
                            255, 0, 0, 255)), QBrush(
                        QColor(
                            255, 0, 0, 100))))
            box_items[item].setAcceptHoverEvents(True)
            box_items[item].setToolTip("Box[%d]: confidence:%s, text:%s" %
                                       (item, conf, result_text))
        QCoreApplication.processEvents()
        self.box_data = box_items

    @pyqtSlot()
    def on_pushButtonLoad_pressed(self):
        """Load Image
        """
        input_dir = sett.readSetting('images/input_dir')
        image = QFileDialog.getOpenFileName(
            self,
            'Open Image file',
            input_dir,
            'Images (*.jpg *.jpeg *.bmp *.png *.tiff *.tif *.gif);;'
            'All files (*.*)')

        if not image:
            self.show_msg('File was not selected…')
            return
        if isinstance(image, tuple):
            filename = image[0]
        else:
            filename = image
        self.load_image(filename)
        sett.storeSetting('images/input_dir', os.path.dirname(filename))

    @pyqtSlot()
    def on_pushButtonRestart_pressed(self):
        """Restart program
        """
        python = sys.executable
        os.execl(python, python, * sys.argv)

    @pyqtSlot()
    def on_actionCrop_triggered(self):
        """Crop image trigger
        """
        self.show_msg("Started cropping: adjust area")
        self.area = AreaItem(0, 0, 300, 150)
        self.area.areaColor = QColor(0, 155, 0, 75)
        self.area.selectedChange.connect(self.areaItemChanges)
        self.scene.addItem(self.area)

    @pyqtSlot('QString')
    def areaItemChanges(self, itemAction):
        if itemAction == 'removeArea':
            self.show_msg("Removing area")
            self.removeSceneItem()
        elif itemAction == 'crop2Area':
            self.crop2Area()

    def removeSceneItem(self):
        """
        Remove items from scene
        """
        items = self.scene.selectedItems()
        if len(items):
            while items:
                item = items.pop()
                self.scene.removeItem(item)
                del item

    def crop2Area(self):
        """Perform crop action on Area item
        """
        xb = int(self.area.pos().x())
        yb = int(self.area.pos().y())
        wb = int(self.area.rect().width())
        hb = int(self.area.rect().height())
        self.show_msg("Cropping area: {}, {}, {}, {}".format(xb, yb, wb, hb))
        box = self.leptonica.boxCreate(xb, yb, wb, hb)
        cropped = self.leptonica.pixClipRectangle(self.pix_image, box, None)
        qimage = lept.pix_to_qimage(self.leptonica, cropped)
        pixmap = QPixmap(qimage)
        if pixmap:
            self.scene.removeItem(self.image)
            del self.image
            self.image = self.scene.addPixmap(pixmap)
            self.scene.setSceneRect(QRectF(pixmap.rect()))
            self.pix_image = cropped
            self.image_width = self.leptonica.pixGetWidth(self.pix_image)
            self.image_height = self.leptonica.pixGetHeight(self.pix_image)
            self.removeSceneItem()
            self.show_msg("New image size: %dx%d, resolution %dx%d" %
                          (self.image_width,
                           self.image_height,
                           self.leptonica.pixGetXRes(self.pix_image),
                           self.leptonica.pixGetYRes(self.pix_image)
                           ))

    @pyqtSlot('QString')
    def load_image(self, filename):
        """Load image to scene and create PIX
        """
        self.image_name = filename
        self.scene.clear()
        self.zoomTo1()
        # Read image with leptonica => create PIX structure and report image
        # size info
        self.pix_image = self.leptonica.pixRead(self.image_name.encode())
        self.image_width = self.leptonica.pixGetWidth(self.pix_image)
        self.image_height = self.leptonica.pixGetHeight(self.pix_image)
        self.show_msg("image size: %dx%d, resolution %dx%d" %
                      (self.image_width,
                       self.image_height,
                       self.leptonica.pixGetXRes(self.pix_image),
                       self.leptonica.pixGetYRes(self.pix_image)
                       ))
        self.box_data = []
        qimage = lept.pix_to_qimage(self.leptonica, self.pix_image)
        if not qimage:  # fallback solution
            pixmap = QPixmap(filename)
        else:
            pixmap = QPixmap(qimage)
        self.image = self.scene.addPixmap(pixmap)
        self.scene.setSceneRect(QRectF(pixmap.rect()))
        self.setWindowTitle(
            'Analyze & ORC image with tesseract and '
            'leptonica :: %s' %
            os.path.basename(
                self.image_name))

    def show_msg(self, message):
        """Show message in textBrowser
        """
        self.textEdit.append(message)
        # Scroll to end of the last message
        cursor = QTextCursor(self.textEdit.textCursor())
        cursor.movePosition(QTextCursor.End)
        self.textEdit.setTextCursor(cursor)
        QApplication.processEvents()

    def closeEvent(self, event):
        """Store setting on exit
        """
        sett.storeSetting('geometry', self.saveGeometry())
        sett.storeSetting('state', self.saveState())
        sett.storeSetting("splitter_1Sizes", self.splitter_1.saveState())
        sett.storeSetting("splitter_1Geo", self.splitter_1.saveGeometry())
        sett.storeSetting("splitter_2Sizes", self.splitter_2.saveState())
        sett.storeSetting('images/last_filename', self.image_name)
        sett.storeSetting('images/zoom_factor', self.getZoomFactor())
        row_l = self.comboBoxLang.currentIndex()
        lang = self.comboBoxLang.itemData(row_l)
        if lang:
            sett.storeSetting('language', lang)
        row_p = self.comboBoxPSM.currentIndex()
        sett.storeSetting('PSM', self.comboBoxPSM.itemData(row_p))
        row_r = self.comboBoxRIL.currentIndex()
        sett.storeSetting('RIL', self.comboBoxRIL.itemData(row_r))
        if self.api:
            self.tesseract.TessBaseAPIEnd(self.api)
        QMainWindow.closeEvent(self, event)

    def setZoom(self, scale):
        """Scale to selected factor
        """
        transform = QTransform()
        transform.scale(scale, scale)
        self.graphicsView.setTransform(transform)

    def zoomFit(self):
        """Zoom image to fit in graphicsView window
        """
        # TODO: put border to preferencies
        border = 0.95
        if (not self.image_height) or (not self.image_width):
            return
        viewWidth = self.graphicsView.viewport().width() * border
        viewHeight = self.graphicsView.viewport().height() * border

        ratio = float(viewWidth / viewHeight)
        aspectRatio = float(self.image_width / self.image_height)

        if ratio > aspectRatio:
            zoomFactor = float(viewHeight / self.image_height)
        else:
            zoomFactor = float(viewWidth / self.image_width)
        self.setZoom(zoomFactor)

    def zoomToWidth(self):
        """Zoom image to fit in graphicsView window
        """
        if not self.image_width:
            return
        viewWidth = self.graphicsView.viewport().width()
        zoomFactor = float(viewWidth / self.image_width)
        self.setZoom(zoomFactor)

    def zoomTo1(self):
        """Zoom to 1:1
        """
        self.setZoom(1)

    def zoomIn(self):
        """Zoom In
        """
        factor = 1.25
        self.graphicsView.scale(factor, factor)

    def zoomOut(self):
        """Zoom Out
        """
        factor = .8
        self.graphicsView.scale(factor, factor)
        self.getZoomFactor()

    def getZoomFactor(self):
        """Get current zoom factor
        """
        return self.graphicsView.transform().m11()

    def eventFilter(self, obj, event):
        """Zoom In/Out with CTRL + mouse wheel
        """
        if event.type() == QEvent.GraphicsSceneWheel:
            assert isinstance(obj, QGraphicsScene)
            delta = event.delta()
            if (event.modifiers() == Qt.ControlModifier and delta > 0):
                factor = 1.41 ** (event.delta() / 240.0)
                self.graphicsView.scale(factor, factor)
                event.accept()
                return True
            elif (event.modifiers() == Qt.ControlModifier and delta < 0):
                factor = 1.41 ** (event.delta() / 240.0)
                self.graphicsView.scale(factor, factor)
                event.accept()
                return True
        return obj.eventFilter(obj, event)


def main():
    """Start GUI
    """
    app = QApplication([])
    if sys.platform == "win32":
        myappid = 'tesseract.python.demo.{}'.format(__version__)  # arbitrary string
        ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(myappid)
    app.setWindowIcon(QIcon(':/images/appIcon'))
    window = MainWindow()
    window.show()
    app.exec_()

if __name__ == '__main__':
    main()
