import codecs, glob, copy, itertools
import xml.etree.ElementTree as ET
import numpy as np
from gensim.models.keyedvectors import KeyedVectors
import sys
from settings import *
from collections import defaultdict


# Container class for VSMs
class VSM:
    def __init__(self, src, vsm_id, pos_aware=False, verbose=False,
                 default=-5.0, mtype="w2v"):
        self.prefix = ""
        if mtype == "w2v":
            self.model = KeyedVectors.load_word2vec_format(src, binary=False)
        else:
            raise Exception("Unsupported model format {}".format(mtype))

        if verbose:
            sys.stderr.write("-> Read VSM: {};\nvocab: {}; dim: {}; pos_aware: {}; default sim: {}\n"
                         .format(src, len(self.model.wv.vocab), self.model.vector_size, pos_aware, default))

        self.pos_aware = pos_aware
        self.default = default  # default similarity for unknown items
        self.vsm_id = vsm_id

        self.src = src

    def prep(self, w, pos=None):
        if self.pos_aware and pos is None:
            raise Exception("The VSM is POS-aware, but no POS provided")

        w = w.lower()  # lowercase
        w = self.prefix + w

        if self.pos_aware:
            w += "|" + pos

        return w

    def vec(self, w, pos=None):
        w = self.prep(w, pos)
        if w in self.model.vocab:
            return self.model[w]
        return None

    def similarity(self, w1, w2, pos1=None, pos2=None):
        if not self.has_word(w1, pos1) or not self.has_word(w2, pos2):
            return self.default

        w1, w2 = self.prep(w1, pos1), self.prep(w2, pos2)

        if w1 not in self.model.wv.vocab or w2 not in self.model.wv.vocab:
            sim = self.default
        else:
            sim = self.model.similarity(w1, w2)
        return sim

    def has_word(self, w, pos=None):
        return self.prep(w, pos) in self.model.wv.vocab

    def vocab(self):
        return self.model.wv.vocab


# Container class for similarity benchmark query sets
class QuerySet:
    FORMATS = ["ws353", "simlex999"]

    def __init__(self, src, data_format, has_header=True, default_pos="N", pos_filter=None):
        if data_format not in self.FORMATS:
            raise Exception("Unsupported format '{}', select from [{}]".format(data_format, "|".join(self.FORMATS)))
        self.queries = {}
        pos_map = {"N": "N", "V": "V", "A": "J"}
        with codecs.open(src, "r", "utf-8") as f:
            for line in f:
                if has_header:
                    has_header = False
                else:
                    if data_format == "ws353":
                        v = line.strip().split("\t")
                        w1, w2, score = v[:3]
                        pos = default_pos
                    elif data_format == "simlex999":
                        v = line.strip().split("\t")
                        w1, w2, pos, score = v[:4]
                        pos = pos_map[pos]
                    if pos_filter is None or pos in pos_filter:
                        self.queries[(w1, w2, pos)] = float(score)

    # compute coverage for the query set given a VSM
    def coverage(self, vsm, pos_filter=None):
        na = []
        total = []
        for (w1, w2, pos) in self.queries:
            if pos_filter is None or pos == pos_filter:
                total += [w1]
                total += [w2]
                if not vsm.has_word(w1, pos):
                    na += [w1]
                if not vsm.has_word(w2, pos):
                    na += [w2]
        total = list(set(total))
        na = list(set(na))
        return len(na), len(total), na


# Generic lexicon container
class Lexicon:
    def __init__(self, src, verbose=False):
        self.lex_map = {}
        self.src = src
        self._read(src)

        if verbose:
            sys.stderr.write("-> Read lexicon: {}\n".format(src))
            sys.stderr.write("items: {}; ".format(len(self.all_members)))
            sys.stderr.write("classes: {}\n".format(len(self.lex_map)))

    def _read(self, src):
        raise NotImplementedError("Use specific lexicon class, e.g. VerbNetLite or WordNetGroupings")

    def members(self, cls=None):
        if cls is None:
            return self.all_members
        if cls not in self.lex_map:
            raise Exception("Unknown class '{}'".format(cls))
        return self.lex_map[cls]

    def classes(self, member=None):
        if member is None:
            return self.all_classes
        return [c for c in self.lex_map if member in self.lex_map[c]]

    def coverage(self, vsm, pos):
        na = []
        ok = []
        for v in self.all_members:
            if not (vsm.has_word(v, pos)):
                na += [v]
            else:
                ok += [v]

        na = list(set(na))
        ok = list(set(ok))
        return len(na), len(self.all_members), na, ok

    def dump(self):
        for attr in dir(self):
            if hasattr(self, attr):
                print("obj.%s = %s" % (attr, getattr(self, attr)))

    def _index(self):
        self.all_members = sorted(list(set(itertools.chain(*self.lex_map.values()))))
        self.all_classes = sorted(list(self.lex_map.keys()))

        self.class_matrix = np.zeros((len(self.all_classes), len(self.all_members)))

        m_ix = {y: x for x, y in enumerate(self.all_members)}

        for c in self.lex_map:
            ic = self.all_classes.index(c)
            for m in self.lex_map[c]:
                im = m_ix[m]
                self.class_matrix[ic][im] = 1.0

        self.class_matrix = self.class_matrix.T


class VerbNetLite(Lexicon):
    def _read(self, src):
        self.lex_map = {}
        for fn in glob.glob(os.path.join(src, "*.xml")):
            vncls = os.path.split(fn)[-1].replace(".xml", "")
            xml = ET.parse(fn)
            members = [x.get("name") for x in xml.findall('.//MEMBER')]
            if len(members) > 0:
                self.lex_map[vncls] = members
        self._index()


class WordNetGroupings(Lexicon):
    def _read(self, src):
        lexref = {}
        self.lex_map = defaultdict(list)
        with codecs.open(WORDNET_LF_REFERENCE, "r", "utf-8") as f:
            for line in f:
                line = line.strip().split()
                lexref[line[0]] = line[1]

            with codecs.open(src, "r", "utf-8") as f:
                for line in f:
                    if not line.startswith(" "):
                        line = line.strip().split()
                        sup = line[1]
                        sup = lexref[sup]

                        lemma_cnt = int(line[3], 16)  # lemma count is a hex ;D
                        lemmas = []
                        for i in range(0, lemma_cnt):
                            lemmas += [line[4+i*2]]  # not the most convenient format

                        lemmas = list(set([l.lower() for l in lemmas if "_" not in l]))
                        self.lex_map[sup] += lemmas
        self.lex_map = {x: sorted(list(set(y))) for x, y in self.lex_map.items()}
        self._index()