'''
How can I load .vec file (pretrained or generated by Fasttext or GloVe/other softwares) to Fasttext to actually use it? Is it possible?
A: it is impossible! You can not recovery a fastText model from .vec files, either.

So, we can use dict() for word2vec and use  KDTree or BallTree for similar words searching.
Also, I implemented class MyNNAlg that plays like KDTree's role
'''

# coding: UTF-8
import fasttext
import mxnet.gluon as gluon
import jieba
import torch
import numpy as np
import random
import torch.utils.data.dataset as dataset
import torch.utils.data.dataloader as dataloader
import torch.nn as nn
import torch.nn.functional as F
import torch.jit
from io import open
import glob
import os
import math
import hashlib
import pickle
import re
from  sklearn.neighbors import KDTree
from  sklearn.neighbors import BallTree
from  sklearn.cluster import KMeans


def cosSim(v1:np.ndarray,v2:np.ndarray):
    return np.dot(v1,v2)/(np.linalg.norm(v1,ord=2)*np.linalg.norm(v2, ord=2))
class MyNNAlg:
    def __init__(self, data:np.ndarray, clusters:int):
        if data.shape[0] > 50000:
            #随机抽取10分之一的数据，用于训练
            sep = data.shape[0]//10
            trainSet = np.zeros((sep, data.shape[1]))
            for i in range(sep):
                r = random.randint(0, data.shape[0])
                trainSet[i] = data[r]
        else:
            sep = data.shape[0]
            trainSet = np.zeros((sep, data.shape[1]))
            for i in range(sep):
                r = random.randint(0, data.shape[0])
                trainSet[i] = data[r]

        self.km = KMeans(n_clusters=clusters, init='random')
        self.km.fit(trainSet)
        print("kmeans fitting finished!")
        self.clusters = [[] for i in range(clusters)]
        for i in range(data.shape[0]):
            c = self.km.predict([data[i]])
            self.clusters[c[0]].append( (i, data[i]))
        print("each sample got its cluster!")

    def query(self, v:np.ndarray, k:int):
        c = self.km.predict(v)
        neighbors = self.clusters[c[0]]
        result = list()
        for id,vec in neighbors:
            sim = cosSim(v[0], vec)
            result.append((sim, id))
        result.sort(key=lambda x:x[0], reverse=True)
        if k < len(result):
            return np.array([[result[i][0] for i in range(k)]]), np.array([[result[i][1] for i in range(k)]])
        else:
            return np.array([[r[0] for r in result]]), np.array([[r[1] for r in result]])


class WordVecSearch:
    def __init__(self, vecFilePath:str):
        self.word2Vec = dict()
        self.word2id = dict()
        self.wordList = list()
        with open(vecFilePath, "r", encoding="utf8") as f:
            cnt = 0
            for line in f:
                if line[-1] == "\n":
                    line = line[:-1]
                if line[-1] == "\r":
                    line = line[:-1]
                cnt += 1
                if cnt == 1:
                    fields = line.split(" ")
                    self.wordCount = int(fields[0])
                    self.dim = int(fields[1])
                    allWordVec = np.zeros((self.wordCount, self.dim), dtype="float32")
                    continue
                fields = line.split(" ")

                if len(fields) < (self.dim+1):
                    raise Exception("invalid line in vec file!" + str(len(fields)))
                for i in range(self.dim):
                    allWordVec[cnt-2, i] = float(fields[i+1])
                self.word2Vec[fields[0]] = allWordVec[cnt-2]
                self.wordList.append(fields[0])
                self.word2id[fields[0]] = len(self.wordList)-1

        #self.tree = KDTree(allWordVec,) # type:sklearn.neighbors.KDtree
        #self.tree = BallTree(allWordVec)
        self.tree = MyNNAlg(allWordVec, int(allWordVec.shape[0]/1000) )
        del allWordVec



    def get_nearest_neighbors(self, word, k=10):
        if word in self.word2Vec:
            v = self.word2Vec[word]
            dist, index = self.tree.query([v], k+1)
            if len(index[0]) <= 1:
                return []
            words = [self.wordList[index[0, i+1]] for i in range(len(index[0])-1)]
            result =  [(cosSim(v, self.word2Vec[w]), w) for w in words]
            result.sort(key=lambda x: x[0], reverse=True)  # 按cosSim 值排序
            return result
        else:
            return []
    def get_word_vector(self, word):
        if word in self.word2Vec:
            v = self.word2Vec[word]
            return  v
        else:
            raise Exception("no this word in vocabulary!"+word)
    def get_word_id(self, word):
        if word in self.word2id:
            return self.word2id[word]
        else:
            return -1
    def get_word(self, id):
        return self.wordList[id]
    def get_words(self):
        return self.wordList
    def get_analogies(self, a, b,c, k=10):
        if a in self.word2Vec and b in self.word2Vec and c in self.word2Vec:
            v = self.word2Vec[a] - self.word2Vec[b] + self.word2Vec[c]
            dist, index = self.tree.query([v], k)
            if len(index[0]) < 1:
                return []
            words = [self.wordList[index[0, i]] for i in range(len(index[0]))]
            result = list()
            for w in words:
                if w == a or w == b or w ==c:
                    continue
                result.append((cosSim(v, self.word2Vec[w]), w))
            result.sort(key=lambda x:x[0], reverse=True) # 按cosSim 值排序
            return result
        else:
            return []


search = WordVecSearch("E:\\DeepLearning\\data\\fastText_model\\wiki.simple.vec")
print(search.get_nearest_neighbors("jlkdsa"))
print(search.get_nearest_neighbors("father"))
print(search.get_analogies("son", "father", "mother"))
print(search.get_analogies("berlin", "germany", "france"))
print(search.get_word_id("father"))
print(search.get_word(13))
print(len(search.get_words()))










