#!/usr/bin/env python
#-*- coding: UTF-8 -*-

__author__ = 'helljump'


# TODO: контекст 3-5 слов в рерайтере


from PyQt4 import QtGui
from PyQt4 import QtCore
import logging
import startup
from ZODB.FileStorage import FileStorage
from ZODB.DB import DB as ZDB
import os
import re
import sys
from os.path import dirname
from ui.rewriter_ui import Ui_Dialog
import pymorphy2
from difflib import SequenceMatcher


log = logging.getLogger(__name__)


def copy_case(text, example):
    k = re.split('\s', example)
    if k[0].istitle():
        text = text.capitalize()
    elif k[0].isupper():
        text = text.upper()
    else:
        text = text.lower()
    return text


def tokenize(text):
    tag_re = re.compile('(</?[^>]+>)', re.U | re.S | re.M)
    word_re = re.compile('([^\w]+)', re.U | re.S | re.M)
    for chunk in re.split(tag_re, text, re.U):
        if chunk and chunk[0] == '<':
            yield chunk
            continue
        for word in re.split(word_re, chunk, re.U):
            yield word


class Rewriter(object):

    def __init__(self):
        if hasattr(sys, "frozen"):
            root = dirname(sys.executable)
        else:
            root = r'build\exe.win32-2.7'  # dirname(__file__)
        root = unicode(root, encoding=sys.getfilesystemencoding())
        self.morph = pymorphy2.MorphAnalyzer(root + '/morph')
        self.st = FileStorage(os.path.join(root + '/syno/db'), read_only=True)
        self.db = ZDB(self.st)
        self.co = self.db.open()
        root = self.co.root()
        self.bdb = root['thes']
        #self.arr = ['syn', 'sim', 'ant', 'lin']

    def close(self):
        self.co.close()
        self.db.close()
        self.st.close()

    def variants(self, word):
        assert isinstance(word, unicode)
        egg = self.morph.parse(word.lower())
        if not egg:
            return
        src_form = egg[0].tag.grammemes
        src_word = egg[0].normal_form.upper()
        record = self.bdb.get(src_word, None)
        if record is None:
            return
        vars = {}
        for k in record.keys():
            vars[k] = []
            for v in record[k]:
                info = self.morph.parse(v)
                if not info:
                    continue
                dst_word = info[0]
                dst_info = dst_word.inflect(src_form)
                if dst_info is None:
                    continue
                rc = copy_case(dst_info.word, word)
                vars[k].append(rc)
        return vars


class Model(QtCore.QAbstractTableModel):

    t = {0: 'syn', 1: 'sim', 2: 'lin', 3: 'ant'}

    def __init__(self, parent=None):
        super(Model, self).__init__(parent)
        self.labels = [u'Синоним', u'Сходный термин', u'Связанный термин', u'Антоним']
        self.data = {}

    def rowCount(self, parent=QtCore.QModelIndex()):
        egg = [len(v) for k, v in self.data.items()]
        return max(egg) if egg else 0

    def columnCount(self, parent=QtCore.QModelIndex()):
        return len(self.labels)

    def data(self, index, role=QtCore.Qt.DisplayRole):
        rc = QtCore.QVariant()
        if not index.isValid():
            return rc
        if not self.data:
            return rc
        row = index.row()
        col = index.column()
        if role == QtCore.Qt.DisplayRole:
            datacol = self.data[self.t[col]]
            if len(datacol) > row:
                return datacol[row]
        return rc

    def flags(self, index):
        egg = QtCore.Qt.NoItemFlags
        row = index.row()
        col = index.column()
        if self.data and len(self.data[self.t[col]]) > row:
            egg |= QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
        return egg

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        if orientation == QtCore.Qt.Horizontal:
            if role == QtCore.Qt.DisplayRole:
                return QtCore.QVariant(self.labels[section])
        return QtCore.QVariant()

    def set_data(self, data):
        if data is None:
            self.data = {}
        else:
            self.data = data
        self.layoutChanged.emit()


class Dialog(QtGui.QDialog, Ui_Dialog):

    def __init__(self, text, parent=None):
        super(Dialog, self).__init__(parent)
        self.setupUi(self)
        self.model = Model(self)
        self.vars_tb.setModel(self.model)
        for i in range(len(self.model.labels)):
            self.vars_tb.horizontalHeader().setResizeMode(i, QtGui.QHeaderView.Stretch)
        self.rw = Rewriter()
        self.tokens = None
        self.text = text
        self.data = []

    def set_next(self):
        if self.tokens is None:
            self.tokens = tokenize(self.text)
        try:
            while True:
                if not self.isVisible():
                    raise StopIteration
                term = self.tokens.next()
                self.termin_le.setText(term)
                QtGui.qApp.processEvents()
                if term and term[0] == '<':
                    self.data.append(term)
                    continue
                if len(term.strip()) < 2:
                    self.data.append(term)
                    continue
                break
            rc = self.rw.variants(term)
            self.model.set_data(rc)
        except StopIteration:
            self.new_text = ''.join(self.data)
            s = SequenceMatcher(None, self.text, self.new_text)
            self.ratio = s.quick_ratio()
            self.hide()
            self.setResult(1)

    @QtCore.pyqtSlot()
    def on_next_pb_clicked(self):
        log.debug('next')
        self.next_pb.setText(u'Заменить')
        self.data.append(unicode(self.termin_le.text()))
        self.set_next()

    @QtCore.pyqtSlot(QtCore.QModelIndex)
    def on_vars_tb_clicked(self, index):
        egg = index.data().toString()
        if not egg.isEmpty():
            self.termin_le.setText(egg)


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    root = startup.CONN.root()
    prj = root["projects"][0]
    prj.open()
    docs = prj.root()['documents']
    dlg = Dialog(docs[0].text)
    rc = dlg.exec_()
    if rc:
        print dlg.ratio
    prj.close()
