"""
Created Thursday June 11 2020 16:32 +0700

@author: arunwpm
"""
try:
    import importlib.resources as pkg_resources
except ImportError:
    import importlib_resources as pkg_resources
from . import data

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import load_model

import traceback
import csv
import pickle
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer

def loadcsv(filename):
    with pkg_resources.open_text(data, filename) as f: 
        return list(csv.reader(f))

model = None
cntVecr = None
tfmer = None
ld = None
id2tag = {}

def initialize( modelfile="model_2500_500_50.h5", 
                vocabfile="small_vocab_20.csv",
                tfmerfile="tfmer_20.p",
                ldloc = 'labelsdict_20.csv', #name of the labels dictionary generated by nb.py (should be labels_dict.csv)
                id2tagloc = 'nyt-theme-tags.csv' #name of the conversion table from tag id to tag name for NYTcorpus
                ):
    global model
    global cntVecr
    global tfmer
    global ld
    global id2tag

    print("Initializing...")
    # initialize the trained model
    model = load_model(modelfile)
    print("Model...")
    
    # initialize the count vectorizer
    vocab = loadcsv(vocabfile)
    vocab = [x[0] for x in vocab] # extra layer of list that we don't want here
    cntVecr = CountVectorizer(vocabulary=vocab)
    print("Count Vectorizer...")

    # initialize the tfidf transformer
    with pkg_resources.open_binary(data, tfmerfile) as tfmerin:
        tfmer = pickle.load(tfmerin)
    print("TF-IDF Transformer...")
    
    # initialize the matrix index -> tag id file and the tag id -> tag name file
    ld = loadcsv(ldloc)
    id2tag_table = loadcsv(id2tagloc)
    for row in id2tag_table:
        if row == []:
            continue
        id2tag[row[1]] = row[2]
    print("Miscellaneous...")

def gettags(txt):
    if (model is None):
        initialize()
    vec0 = cntVecr.transform([txt])
    vec1 = tfmer.transform(vec0)
    print(vec1)

    mat = model.predict(vec1.todense())
    print(mat)

    tags = []
    for i in range(len(mat[0])):
        if float(mat[0][i]) >= 0.5:
            tags.append(id2tag[ld[i][1]])
    print(tags)

    return tags

if __name__ == "__main__":
    while True:
        txt = input("Enter text: ")
        gettags(txt)