"""Build the data embedded in the word-map simulator of the Word Embeddings page.

NOT a figure script (the leading underscore keeps generate-figures.sh away):
it needs gensim, which the project venv does not have. Run it with a Python
that has gensim + scikit-learn, from this folder:

    pip install gensim scikit-learn
    python _build_sim_data.py

It loads the pretrained GloVe vectors (glove-wiki-gigaword-50: 400,000 words,
50 dimensions), keeps the words listed below, and writes into
docs/assets/javascripts/sims/sim-word-embeddings.js, between the /*DATA*/ and
/*END-DATA*/ markers:
  - each word's unit-length 50-d vector, rounded to 3 decimals (stored x1000);
  - two 2-D layouts of those vectors: PCA and t-SNE (scikit-learn);
  - the PCA mean and axes, so the simulator can project a computed vector;
  - for each preset analogy, the answer searched over ALL 400,000 GloVe words.
"""
import json
import re
from pathlib import Path

import gensim.downloader as api
import numpy as np
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE

GROUPS = [
    (('people', 'pessoas'), 'man woman king queen prince princess boy girl father mother son daughter '
                            'brother sister husband wife uncle aunt he she'),
    (('countries & capitals', 'países e capitais'), 'france italy germany spain japan china russia brazil portugal '
                                                    'egypt greece canada paris rome berlin madrid tokyo beijing '
                                                    'moscow brasilia lisbon cairo athens ottawa'),
    (('verb & adjective forms', 'formas de verbos e adjetivos'), 'walk walked swim swam run ran go went eat ate '
                                                                 'write wrote good better bad worse big bigger '
                                                                 'small smaller fast faster slow slower'),
    (('animals', 'animais'), 'cat dog horse cow lion tiger puppy kitten bird fish'),
    (('food', 'comida'), 'apple banana mango orange grape bread cheese rice coffee tea'),
    (('transport', 'transporte'), 'car truck bus train plane bicycle boat ship'),
    (('occupations', 'profissões'), 'doctor nurse engineer teacher programmer receptionist homemaker carpenter '
                                    'lawyer scientist captain maid'),
    (('technology', 'tecnologia'), 'computer software internet microsoft google iphone python java'),
    (('money & rivers', 'dinheiro e rios'), 'bank money loan credit stock market river water shore lake'),
    (('nature & colors', 'natureza e cores'), 'hot cold dark light sun moon day night red blue green yellow'),
    (('music & sport', 'música e esporte'), 'music guitar piano football soccer'),
]
# a - b + c  ("a is to b as ? is to c", written the way the simulator shows it)
PRESETS = [('king', 'man', 'woman'), ('tokyo', 'japan', 'france'), ('swam', 'swim', 'walk'),
           ('berlin', 'germany', 'italy'), ('bigger', 'big', 'small'), ('brother', 'man', 'woman'),
           ('doctor', 'man', 'woman'), ('night', 'moon', 'sun')]

kv = api.load('glove-wiki-gigaword-50')
words, group = [], []
for g, (_, ws) in enumerate(GROUPS):
    for w in ws.split():
        assert w in kv.key_to_index, w
        words.append(w)
        group.append(g)
V = np.array([kv.get_vector(w, norm=True) for w in words])
Vq = np.round(V * 1000).astype(int)                 # what the simulator stores
Vr = Vq / 1000.0
Vr /= np.linalg.norm(Vr, axis=1, keepdims=True)

pca = PCA(n_components=2).fit(Vr)
P = pca.transform(Vr)
T = TSNE(n_components=2, perplexity=25, metric='cosine', init='pca', random_state=0).fit_transform(Vr)
print(f"{len(words)} words, {V.shape[1]} dimensions")
print("PCA: variance kept by the 2 axes:", pca.explained_variance_ratio_.round(3),
      f"total {pca.explained_variance_ratio_.sum():.1%}")
print("largest cosine change caused by rounding to 3 decimals:",
      f"{np.abs(V @ V.T - Vr @ Vr.T).max():.4f}")

presets = []
for a, b, c in PRESETS:
    full = kv.most_similar(positive=[a, c], negative=[b], topn=3)
    presets.append({'a': a, 'b': b, 'c': c, 'full': [[w, round(float(s), 2)] for w, s in full]})
    print(f"{a} - {b} + {c} -> all 400k words: {full[0][0]} {full[0][1]:.2f}")

data = {
    'words': words, 'g': group, 'groups': [list(n) for n, _ in GROUPS], 'dim': int(V.shape[1]),
    'vec': Vq.ravel().tolist(),
    'pca': np.round(P, 3).ravel().tolist(), 'tsne': np.round(T, 2).ravel().tolist(),
    'pcaMean': np.round(pca.mean_, 5).tolist(), 'pcaComp': np.round(pca.components_, 5).tolist(),
    'pcaVar': np.round(pca.explained_variance_ratio_, 4).tolist(), 'presets': presets,
}
js = Path(__file__).resolve().parents[5] / 'docs/assets/javascripts/sims/sim-word-embeddings.js'
src = js.read_text(encoding='utf-8')
blob = json.dumps(data, separators=(',', ':'), ensure_ascii=False)
src = re.sub(r'/\*DATA\*/.*?/\*END-DATA\*/', lambda m: '/*DATA*/' + blob + '/*END-DATA*/', src, flags=re.S)
js.write_text(src, encoding='utf-8')
print(f"wrote {len(blob) / 1024:.1f} KB of data into {js.name}")