from pygoose import *

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_distances, euclidean_distances

"""
Explanation: Feature: TF-IDF Distances
Create TF-IDF vectors from question texts and compute vector distances between them.
Imports
This utility package imports numpy, pandas, matplotlib and a helper kg module into the root namespace.
End of explanation
"""


project = kg.Project.discover()

"""
Explanation: Config
Automatically discover the paths to various data folders and compose the project structure.
End of explanation
"""


feature_list_id = 'tfidf'

"""
Explanation: Identifier for storing these features on disk and referring to them later.
End of explanation
"""


tokens_train = kg.io.load(project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_no_stopwords_train.pickle')
tokens_test = kg.io.load(project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_no_stopwords_test.pickle')

tokens = tokens_train + tokens_test

"""
Explanation: Read Data
Preprocessed and tokenized questions.
End of explanation
"""


all_questions_flat = np.array(tokens).ravel()

documents = list(set(' '.join(question) for question in all_questions_flat))

del all_questions_flat

"""
Explanation: Extract a set of unique question texts (document corpus).
End of explanation
"""


vectorizer = TfidfVectorizer(
    encoding='utf-8',
    analyzer='word',
    strip_accents='unicode',
    ngram_range=(1, 1),
    lowercase=True,
    norm='l2',
    use_idf=True,
    smooth_idf=True,
    sublinear_tf=True,
)

vectorizer.fit(documents)

model_filename = 'tfidf_vectorizer_{}_ngrams_{}_{}_penalty_{}.pickle'.format(
    vectorizer.analyzer,
    vectorizer.ngram_range[0],
    vectorizer.ngram_range[1],
    vectorizer.norm,
)

kg.io.save(vectorizer, project.trained_model_dir + model_filename)

"""
Explanation: Train TF-IDF vectorizer
Create a bag-of-token-unigrams vectorizer.
End of explanation
"""


def compute_pair_distances(pair):
    q1_doc = ' '.join(pair[0])
    q2_doc = ' '.join(pair[1])
    
    pair_dtm = vectorizer.transform([q1_doc, q2_doc])
    q1_doc_vec = pair_dtm[0]
    q2_doc_vec = pair_dtm[1]
    
    return [
        cosine_distances(q1_doc_vec, q2_doc_vec)[0][0],
        euclidean_distances(q1_doc_vec, q2_doc_vec)[0][0],
    ]

features = kg.jobs.map_batch_parallel(
    tokens,
    item_mapper=compute_pair_distances,
    batch_size=1000,
)

X_train = np.array(features[:len(tokens_train)], dtype='float64')
X_test = np.array(features[len(tokens_train):], dtype='float64')

print('X_train:', X_train.shape)
print('X_test: ', X_test.shape)

"""
Explanation: Vectorize train and test sets, compute distances
End of explanation
"""


feature_names = [
    'tfidf_cosine',
    'tfidf_euclidean',
]

project.save_features(X_train, X_test, feature_names, feature_list_id)

"""
Explanation: Save features
End of explanation
"""


