import json
import os
import sys

sys.path.append(os.path.realpath("."))

import multiprocessing

import faiss
import numpy as np
import pandas as pd
from gensim.models import Word2Vec, word2vec
from utils.defs import nth_element
from utils.files import get_df
from utils.cache import redis_client_db


def device_tractate_fe():
    device_tags_df = get_df("personas_device_feature.csv")
    device_tags_df = device_tags_df[["cl_id", "business_tags"]]
    device_tags_df["business_tags"] = device_tags_df["business_tags"].str.split(",").\
        apply(lambda d: d if isinstance(d, list) else [])
    print(device_tags_df.head(3))

    tractate_tags_df = get_df("personas_tractate_tags.csv", columns=["tractate_id", "business_tags"])
    tractate_tags_df["business_tags"] = tractate_tags_df["business_tags"].str.split(",").\
        apply(lambda d: d if isinstance(d, list) else [])
    print(tractate_tags_df.head(3))

    return device_tags_df, tractate_tags_df


def tractate_business_tags_word2vec(data):
    print(len(data))
    model = Word2Vec(data, hs=0, min_count=3, workers=multiprocessing.cpu_count(), iter=10)
    print(model)
    return model


if __name__ == "__main__":
    device_tags_df, tractate_tags_df = device_tractate_fe()
    tags_data = tractate_tags_df["business_tags"].to_list()
    model = tractate_business_tags_word2vec(tags_data)

    # all business tags
    tags_set = set()
    for i in tags_data:
        for j in i:
            tags_set.add(j)

    # tag vector dict
    tags_vector_dict = {}
    for i in tags_set:
        try:
            vec = json.dumps(model.wv.get_vector(i).tolist())
            # vec = model.wv.get_vector(i)
            tags_vector_dict[i] = vec
        except Exception as e:
            pass
    redis_client_db.hmset("personas_tags_embedding", tags_vector_dict)

    print(len(tags_vector_dict.items()))
    # print(random.choice(list(tags_vector_dict.items())))

    # tractate vector dict
    tractate_vector_dict = {}
    for _, row in tractate_tags_df.iterrows():
        vecs = []
        for i in row["business_tags"]:
            vec = tags_vector_dict.get(i)
            if vec:
                vecs.append(np.array(json.loads(vec)).astype("float32"))
        if vecs:
            tractate_vector_dict[row["tractate_id"]] = np.average(vecs, axis=0)
    print(len(tractate_vector_dict.items()))
    # print(random.choice(list(tractate_vector_dict.items())))

    # tractate vector index
    tractate_ids = np.array(list(tractate_vector_dict.keys())).astype("int")
    tractate_embeddings = np.array(list(tractate_vector_dict.values())).astype("float32")

    index = faiss.IndexFlatL2(tractate_embeddings.shape[1])
    print("trained: " + str(index.is_trained))

    index2 = faiss.IndexIDMap(index)
    index2.add_with_ids(tractate_embeddings, tractate_ids)
    print("trained: " + str(index2.is_trained))
    print("total index: " + str(index2.ntotal))

    base_dir = os.getcwd()
    model_dir = os.path.join(base_dir, "_models")
    index_path = os.path.join(model_dir, "faiss_personas_vector.index")
    faiss.write_index(index2, index_path)
    print(index_path)

    # device vector
    # for _, row in device_tags_df.iterrows():
    #     vecs = []
    #     for i in row["business_tags"]:
    #         # vec = tags_vector_dict.get(i, np.array([]))
    #         vec = tags_vector_dict.get(i)
    #         if vec:
    #             vecs.append(np.array(json.loads(vec)).astype("float32"))
    #     if vecs:
    #         t = np.array([np.average(vecs, axis=0)]).astype("float32")
    #         D, I = index2.search(t, 10)
    #         print(row["cl_id"], row["business_tags"])
    #         print(I)

# curl "http://172.16.31.17:9000/gm-dbmw-tractate-read/_search?pretty" -d '{"query": {"term": {"id": "10269"}}, "_source": {"include": ["content", "portrait_tag_name"]}}'
