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

from __future__ import unicode_literals

import array
import logging
import re
import struct

from io import open, StringIO

from unitex import UnitexException, UnitexConstants
from unitex.io import UnitexFile
from unitex.utils.fsa import FSAConstants, Automaton
from unitex.utils.types import BRACKETED_ENTRY, Tag, Entry

_LOGGER = logging.getLogger(__name__)



class CompressedEntry(Entry):

    SEPARATORS = (" ", "-")
    SPLITTER = re.compile("([-\s])")

    def __init__(self):
        super(CompressedEntry, self).__init__()

    def compute(self, lemma, form):
        n, i = "", 0

        while i < len(lemma) and lemma[i].isdigit():
            n = n + lemma[i]
            i = i + 1

        if i > 0:
            prefix = form[:len(form)-int(n)]
        else:
            prefix = form

        suffix = lemma[i:]

        return "%s%s" % (prefix, suffix)

    def uncompress(self, lemma):
        form = self.get_form()
        if not lemma:
            return form

        # If two words don't have de same number of elements
        # the compressed lemma is preceded by '_'
        if lemma[0] == '_':
            return self.compute(lemma[1:], form)

        wtab = self.SPLITTER.split(form)
        ltab = self.SPLITTER.split(lemma)

        l = []
        for i in range(len(ltab)):
            if not ltab[i]:
                continue
            elif ltab[i] in self.SEPARATORS:
                l.append(ltab[i])
            else:
                l.append(self.compute(ltab[i], wtab[i]))

        return "".join(l)

    def load(self, form, data, lemmatize=True):
        data = data.rstrip()

        self.set_form(form)
        lemma = ""

        i = 0

        lemma, escaped = "", False
        try:
            while True:
                if data[i] == "." and escaped is False:
                    break
                elif data[i] == "\\":
                    if escaped is True:
                        lemma += data[i]
                        escaped = False
                    else:
                        lemma += data[i]
                        escaped = True
                else:
                    lemma += data[i]
                    escaped = False
                i += 1
        except IndexError:
            raise UnitexException("Wrong lemma for entry '%s' ..." % data)

        if lemmatize is True:
            self.set_lemma(self.uncompress(lemma))

        Tag.load(self, data[i+1:])



class OldCompressedDictionary:

    INITIAL_STATE_OFFSET=4
    INF_SEPARATOR=re.compile(r"(?<![\\]),")

    def __init__(self):
        self.__bin = None
        self.__inf = None

        self.__buffer = None

    def lookup(self, token, i=None, pos=None):
        if i is None:
            i = 0

        if pos is None:
            pos = self.INITIAL_STATE_OFFSET
        tnbr = self.__bin[pos] * 256 + self.__bin[pos+1]
        pos = pos + 2

        _LOGGER.debug(u"Lookup Start: token[%s|%s] -- pos(%s) -- tnbr(%s)\n" % (token[:i], token[i:], pos, tnbr))

        if i == len(token):
            data = []

            _LOGGER.debug(u"   Check Final State: pos(%s) -- tnbr(%s)\n" % (pos, tnbr))
            if not (tnbr & 32768):
                _LOGGER.debug(u"      -> Final\n")
                index = self.__bin[pos] * 256 * 256 + self.__bin[pos+1] * 256 + self.__bin[pos+2]

                for inf in self.INF_SEPARATOR.split(self.__inf[index]):
                    E = CompressedEntry()
                    E.load(token, inf)

                    data.append(E)
            else:
                _LOGGER.debug(u"      -> Not final\n")

            return data, pos-2
        elif tnbr & 32768:
            tnbr = tnbr - 32768
        else:
            pos = pos + 3

        for j in range(tnbr):
            char = chr(self.__bin[pos] * 256 + self.__bin[pos+1])
            _LOGGER.debug(u"   Matching char[%s] -- pos(%s) -> current[%s]\n" % (token[i], pos, char))

            pos = pos + 2

            offset = self.__bin[pos] * 256 * 256 + self.__bin[pos+1] * 256 + self.__bin[pos+2]
            pos = pos + 3

            if char == token[i]:
                _LOGGER.debug(u"      -> Char found\n")
                return self.lookup(token, i+1, offset)

            # WEIRD... Objective: handle whitespaces in MWU dictionaries for the match function
            #               -> ["Conseil", "d'", "administration"] == "Conseil d'administration"
            elif char == u" " and i == 0:
                _LOGGER.debug(u"   -> Char is whitespace [pass]\n")
                return self.lookup(token, i, offset)

        return None, pos

    def find(self, token):
        entries, pos = self.lookup(token)
        return entries

    def match(self, sequence, i=None, mode=None, separator=None):
        if i is None:
            i = 0

        if mode is None:
            mode = UnitexConstants.MATCH_MODE_LONGEST
        elif mode not in [UnitexConstants.MATCH_MODE_LONGEST,\
                          UnitexConstants.MATCH_MODE_SHORTEST,\
                          UnitexConstants.MATCH_MODE_ALL]:
            raise UnitexException("Wrong match mode: %s ..." % mode)

        matches = []

        buffer, pos, tnbr = [], None, None
        for j in range(i, len(sequence)):
            _LOGGER.debug(u"Match Token: '%s'\n" % sequence[j])

            entries, pos = self.lookup(sequence[j], pos=pos)
            if entries is None:
                _LOGGER.debug(u"   -> No entry found ...\n")
                break
            _LOGGER.debug(u"   -> Entries found: pos[%s] -- tnbr[%s]\n" % (pos, tnbr))

            buffer.append(j)

            if entries:
                matches.append((entries, buffer[:]))
                if mode == UnitexConstants.MATCH_MODE_SHORTEST:
                    return matches

            if separator is not None:
                _LOGGER.debug(u"Match Separator: '%s'\n" % separator)
                entries, pos = self.lookup(separator, pos=pos)
                if entries is None:
                    _LOGGER.debug(u"   -> No separator found ...\n")
                    break
                _LOGGER.debug(u"   -> Separator found\n")

        if not matches:
            return None
        elif mode == UnitexConstants.MATCH_MODE_LONGEST:
            return [matches[-1]]
        elif mode == UnitexConstants.MATCH_MODE_ALL:
            return matches

    def dump(self, pos=None):
        if pos is None:
            pos = self.INITIAL_STATE_OFFSET
            self.__buffer = []

        tnbr = self.__bin[pos] * 256 + self.__bin[pos+1]
        pos = pos + 2

        if not (tnbr & 32768):
            index = self.__bin[pos] * 256 * 256 + self.__bin[pos+1] * 256 + self.__bin[pos+2]

            form = "".join(self.__buffer)

            for inf in self.INF_SEPARATOR.split(self.__inf[index]):
                E = CompressedEntry()
                E.load(form, inf)
                yield E

            pos = pos + 3

        else:
            tnbr = tnbr - 32768

        for j in range(tnbr):
            self.__buffer.append(chr(self.__bin[pos] * 256 + self.__bin[pos+1]))
            pos = pos + 2

            offset = self.__bin[pos] * 256 * 256 + self.__bin[pos+1] * 256 + self.__bin[pos+2]
            pos = pos + 3

            for E in self.dump(offset):
                yield E

        if self.__buffer:
            self.__buffer.pop()

    def load(self, bin, inf, encoding=None):
        if encoding is None:
            encoding = UnitexConstants.DEFAULT_ENCODING
        INF = open(inf, "r", encoding=encoding)

        self.__inf = INF.readlines()
        self.__inf.pop(0) # Remove number information

        INF.close()

        BIN = open(bin, "r+b")

        a = struct.unpack('B', BIN.read(1))[0]
        b = struct.unpack('B', BIN.read(1))[0]
        c = struct.unpack('B', BIN.read(1))[0]
        d = struct.unpack('B', BIN.read(1))[0]
        size = d + (256*c) + (256*256*b) + (256*256*256*a)

        BIN.close()

        BIN = open(bin, "rb")

        self.__bin = array.array('B')

        byte = BIN.read(1)
        while byte:
            tmp = struct.unpack('B', byte)[0]

            self.__bin.append(tmp)

            byte = BIN.read(1)

        BIN.close()



class CompressedDictionary(OldCompressedDictionary):

    def __init__(self):
        super(CompressedDictionary, self).__init__()
        raise NotImplementedError



class GRF(Automaton):

    def __init__(self, name="GRF"):
        super(GRF, self).__init__(name)

    def load(self, file, encoding=None):
        if encoding is None:
            encoding = UnitexConstants.DEFAULT_ENCODING
        raise NotImplementedError

    def save(self, file, encoding=None):
        if encoding is None:
            encoding = UnitexConstants.DEFAULT_ENCODING

        X = 1000
        Y = 1000
        GAP = 20

        transitions = []
        transitions.append({"label": FSAConstants.EPSILON, "targets": set([])})
        transitions.append({"label": "", "targets": set([])})

        nmap = {}
        root = []

        for edge, sid, tid in self.iter("dfs"):
            source = self[sid]
            target = self[tid]

            index = 0

            key = (str(edge), tid)
            if key in nmap:
                index = nmap[key]
            else:
                index = len(transitions)
                nmap[key] = index
                transitions.append({"label": str(edge), "targets": set([])})

            if sid == self.get_initial():
                transitions[0]["targets"].add(str(index))
            if target.is_final() is True:
                transitions[index]["targets"].add("1")

            for _edge in target:
                for _target in target[_edge]:
                    _index = 0

                    _key = (str(_edge), _target.get_id())
                    if _key in nmap:
                        _index = nmap[_key]
                    else:
                        _index = len(transitions)
                        nmap[_key] = _index
                        transitions.append({"label": str(_edge), "targets": set([])})

                    transitions[index]["targets"].add(str(_index))

        with open(file, "w", encoding=encoding) as output:
            output.write("#Unigraph\r\n")
            output.write("SIZE %s %s\r\n" % (X+GAP, Y+GAP))
            output.write("FONT Times New Roman:B 10\r\n")
            output.write("OFONT Monospaced:B 8\r\n")
            output.write("BCOLOR 16777215\r\n")
            output.write("FCOLOR 0\r\n")
            output.write("ACOLOR 13487565\r\n")
            output.write("SCOLOR 16711680\r\n")
            output.write("CCOLOR 255\r\n")
            output.write("DBOXES y\r\n")
            output.write("DFRAME y\r\n")
            output.write("DDATE y\r\n")
            output.write("DFILE y\r\n")
            output.write("DDIR n\r\n")
            output.write("DRIG n\r\n")
            output.write("DRST n\r\n")
            output.write("FITS 100\r\n")
            output.write("PORIENT L\r\n")
            output.write("#\r\n")
            output.write("%s\r\n" % len(transitions))

            for transition in transitions:
                label = transition["label"]
                size = len(transition["targets"])
                targets = " ".join(list(transition["targets"]))

                if size == 0:
                    output.write('"%s" %s %s %s \r\n' % (label, GAP, GAP, size))
                else:
                    output.write('"%s" %s %s %s %s \r\n' % (label, GAP, GAP, size, targets))



class SentenceFST(Automaton):

    def __init__(self, name="SentenceFST"):
        super(SentenceFST, self).__init__(name)

        self.__sentence = None

        self.__tokens = None
        self.__entries = None

    def get_sentence(self):
        return self.__sentence

    def get_tokens(self):
        return self.__tokens

    def get_token(self, i):
        try:
            return self.__tokens[i]
        except IndexError:
            raise UnitexException("SentenceFST token index out of range (size: %s)." %\
                                  len(self.__tokens))

    def get_entries(self):
        return self.__entries

    def get_label(self, key):
        try:
            return self.__entries[key]
        except KeyError:
            raise UnitexException("SentenceFST label key error.")

    def load(self, sentence, tokens, states, entries):
        self.__sentence = sentence

        self.__tokens = []
        self.__entries = {}

        start = 0
        for index, length in tokens:
            end = start + length

            self.__tokens.append(self.__sentence[start:end])
            start = end

        transitions = []

        for i in range(len(states)):
            initial = False
            if i == 0:
                initial = True

            final = False
            if states[i] == "t":
                final = True

            sid = self.add_node(initial=initial, final=final)
            if final is True:
                break

            for lid, tid in states[i]:
                entry = entries[lid][0]

                p1 = entries[lid][1][0][0]
                p2 = entries[lid][1][1][0]

                if p1 not in self.__entries:
                    self.__entries[p1] = []
                self.__entries[p1].append((entry, p2))

                #transitions.append((sid, lid, tid))
                transitions.append((sid, entry, tid))

        #for sid, lid, tid in transitions:
        for sid, entry, tid in transitions:
            self.add_edge(entry, sid, tid)



class TextFST:

    def __init__(self):
        self.__tfst = None
        self.__tind = []

    def __del__(self):
        self.__tfst.close()

    def __len__(self):
        return len(self.__tind)

    def __readline(self, position):
        end = None
        return end

    def __getitem__(self, i):
        try:
            position = self.__tind[i]
        except IndexError:
            raise UnitexException("TextFST index out of range (size: %s)." % len(self))

        self.__tfst.seek(position)

        line = self.__tfst.readline()
        while line:
            line = line.rstrip()

            if line[0] != "$":
                raise UnitexException("Text automaton file is corrupted ...")

            # The sentence number (format '$n')
            number = int(line[1:])

            line = self.__tfst.readline()
            line = line.rstrip()

            # The text of the sentence
            text = line

            line = self.__tfst.readline()
            line = line.rstrip()

            # The tokens of the text
            #   -> [(x1, y), (x2, y2), ..., (xi, yi)]
            #      where,
            #        - x: token index in file 'tokens.txt'
            #        - y: length of the token (in characters)
            tokens = [tuple(int(t) for t in token.split("/")) for token in line.split(" ")]

            line = self.__tfst.readline()
            line = line.rstrip()

            # The offset of the sentence (from the begining of the text)
            #   -> X_Y
            #      where,
            #        - X: the offset in tokens
            #        - Y: the offset in characters
            offset = tuple(int(o) for o in line.split("_"))

            line = self.__tfst.readline()
            line = line.rstrip()

            states = []
            while line != "t":
                if line[0] != ":":
                    raise UnitexException("File '%s' is corrupted ..." % self.__tfst.name)

                line = line[1:].strip()
                line = line.split()

                state = []
                for i in range(0, len(line), 2):
                    state.append((int(line[i]), int(line[i+1])))
                states.append(state)

                line = self.__tfst.readline()
                line = line.rstrip()

                if not line:
                    raise UnitexException("File '%s' is corrupted ..." % self.__tfst.name)

            states.append(line)

            line = self.__tfst.readline()
            line = line.rstrip()

            if line[0] != "f":
                raise UnitexException("File '%s' is corrupted ..." % self.__tfst.name)

            line = self.__tfst.readline()
            line = line.rstrip()

            tags = []
            while line != "f":
                if line == "@<E>":
                    tags.append(("<E>", None))

                elif line == "@STD":
                    line = self.__tfst.readline()
                    line = line.rstrip()

                    content = line[1:]

                    entry = Entry()

                    if BRACKETED_ENTRY.match(content):
                        content = BRACKETED_ENTRY.sub(r"\1", content)
                        entry.load(content)
                    else:
                        entry.set_form(content)

                    line = self.__tfst.readline()
                    line = line.rstrip()

                    if line[0] != "@":
                        raise UnitexException("File '%s' is corrupted ..." % self.__tfst.name)

                    position = [tuple(int(i) for i in p.split(".")) for p in line[1:].split("-")]

                    tags.append((entry, position))

                else:
                    raise UnitexException("File '%s' is corrupted ..." % self.__tfst.name)

                line = self.__tfst.readline()
                line = line.rstrip()

                if line[0] != ".":
                    raise UnitexException("File '%s' is corrupted ..." % self.__tfst.name)

                line = self.__tfst.readline()
                line = line.rstrip()

            _LOGGER.debug(u"SENTENCE[%s]\n" % number)
            _LOGGER.debug(u" - offset: (%s)\n" % ", ".join([unicode(i) for i in offset]))
            _LOGGER.debug(u" - text: %s\n" % text)
            _LOGGER.debug(u" - tokens: [%s]\n" % ", ".join([unicode(t) for t in tokens]))
            _LOGGER.debug(u" - states:\n")
            for state in states:
                _LOGGER.debug(u"   - s: %s\n" % state)
            _LOGGER.debug(u" - tags:\n")
            for tag in tags:
                _LOGGER.debug(u"   - t: (%s)\n" % ", ".join([unicode(t) for t in tag]))

            S = SentenceFST(u"SENTENCE[%d]" % number)
            S.load(text, tokens, states, tags)

            return S

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]

    def load(self, fst, ind, encoding=None):
        if encoding is None:
            encoding = UnitexConstants.DEFAULT_ENCODING

        ffst = UnitexFile()
        ffst.open(fst, "r")
        self.__tfst = StringIO(ffst.read())
        ffst.close()

        self.__tind = []

#        find = UnitexFile()
#        find.open(ind, "b")
#        index = find.read()
#        find.close()
#
#        i = 0
#        while i<(len(index)-4):
#            position = struct.unpack("<L", index[i:i+4])
#            self.__tind.append(position[0])
#        
#            i += 4

        # When using StringIO, the offsets (which are correct in Vim)
        # given by the index file are not correct... WHY ?!
        position = self.__tfst.tell()
        line = self.__tfst.readline()
        while line:
            if line.startswith("$"):
                self.__tind.append(position)
            position = self.__tfst.tell()
            line = self.__tfst.readline()
