"""Build the data embedded in the semantic-search simulator of the Sentence Embeddings page.

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

    pip install sentence-transformers scikit-learn
    python _build_sim_data.py

It encodes the 40 sentences and 10 queries below with the multilingual model
paraphrase-multilingual-MiniLM-L12-v2 (384 numbers per text, unit length),
stores each vector as 384 signed bytes plus one scale (int8 quantization,
base64), and writes everything into
docs/assets/javascripts/sims/sim-sentence-search.js between the /*DATA*/ and
/*END-DATA*/ markers. The TF-IDF ranking is NOT stored: the simulator
computes it in the browser, with the same settings as the TfidfVectorizer
below, which this script uses only to print a reference ranking.
"""
import base64
import json
import re
from pathlib import Path

import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import TfidfVectorizer

MODEL = 'paraphrase-multilingual-MiniLM-L12-v2'
SENTENCES = [
    # 0-4 account
    "Como faço para trocar a senha da minha conta?",
    "I forgot my login credentials.",
    "Não consigo entrar no aplicativo desde ontem.",
    "How do I turn on two-factor authentication?",
    "Meu usuário foi bloqueado depois de três tentativas.",
    # 5-9 delivery
    "Qual é o prazo de entrega para o Nordeste?",
    "Where is my package? It should have arrived by now.",
    "O pedido está parado na transportadora há uma semana.",
    "Can I change the delivery address after ordering?",
    "Vocês fazem entregas aos sábados?",
    # 10-14 payment
    "Quais formas de pagamento vocês aceitam?",
    "Can I pay with Pix?",
    "Meu cartão foi recusado na hora de finalizar a compra.",
    "Erro 402 ao pagar com cartão de crédito.",
    "Erro 404 ao abrir a página do produto.",
    # 15-19 returns and refunds
    "Como devolvo um produto com defeito?",
    "I want a refund for a phone that arrived broken.",
    "O prazo para desistir da compra é de sete dias.",
    "Can I exchange the shoes for a bigger size?",
    "Quanto tempo leva para o estorno aparecer na fatura?",
    # 20-24 bank / bench / river bank
    "O banco aumentou a taxa de juros do financiamento.",
    "The bank approved my loan application.",
    "Sentei no banco da praça para descansar.",
    "We had a picnic on the bank of the river.",
    "Comprei um banco de madeira para a cozinha.",
    # 25-29 weather
    "Vai chover amanhã em São Paulo?",
    "It is freezing outside this morning.",
    "A previsão indica calor de 35 graus no fim de semana.",
    "Leve um guarda-chuva, o céu está muito nublado.",
    "The storm knocked down several trees last night.",
    # 30-34 cooking
    "Receita de bolo de cenoura com cobertura de chocolate.",
    "How long should I boil an egg?",
    "Esse restaurante serve comida japonesa muito boa.",
    "Deixe a massa descansar por trinta minutos antes de assar.",
    "The soup needs more salt and a little pepper.",
    # 35-39 football
    "O Brasil venceu a Argentina nos pênaltis.",
    "The match was postponed because of the rain.",
    "O atacante marcou dois gols no segundo tempo.",
    "Who won the championship last season?",
    "O técnico foi demitido depois de cinco derrotas seguidas.",
]
# query -> ids of the sentences a person would accept as a good answer
QUERIES = [
    ("esqueci minha senha", [0, 1]),
    ("my order never arrived", [6, 7]),
    ("dá para pagar com pix?", [10, 11]),
    ("quero meu dinheiro de volta", [16, 19, 17]),
    ("erro 402", [13]),
    ("um lugar para sentar", [22, 24]),
    ("the bank raised interest rates", [20]),
    ("vai fazer frio amanhã?", [25, 26, 27]),
    ("como cozinhar ovos", [31]),
    ("quem ganhou o jogo?", [35, 38]),
]

model = SentenceTransformer(MODEL)
E = model.encode(SENTENCES, normalize_embeddings=True)
Q = model.encode([q for q, _ in QUERIES], normalize_embeddings=True)


def quantize(M):
    scale = np.abs(M).max(axis=1) / 127.0
    q = np.round(M / scale[:, None]).astype(np.int8)
    return q, scale


qE, sE = quantize(E)
qQ, sQ = quantize(Q)
deq = lambda q, s: (q.astype(float) * s[:, None]) / np.linalg.norm(q.astype(float) * s[:, None], axis=1, keepdims=True)
err = np.abs(Q @ E.T - deq(qQ, sQ) @ deq(qE, sE).T).max()
print(f"{len(SENTENCES)} sentences, {len(QUERIES)} queries, {E.shape[1]} numbers each")
print(f"largest cosine change caused by int8 quantization: {err:.4f}")

# reference rankings: TF-IDF (the simulator recomputes it in JavaScript) and the
# embeddings after int8 quantization, exactly as the simulator ranks them
tf = TfidfVectorizer(token_pattern=r'(?u)\b\w\w+\b').fit(SENTENCES)
T = tf.transform(SENTENCES)
hits = {'TF-IDF': 0, 'embedding': 0}
total = 0
Ed, Qd = deq(qE, sE), deq(qQ, sQ)                  # the vectors the simulator actually uses
for (q, rel), qv in zip(QUERIES, Qd):
    t = (tf.transform([q]) @ T.T).toarray()[0]
    e = Ed @ qv
    top_t, top_e = np.argsort(-t, kind='stable')[:3], np.argsort(-e, kind='stable')[:3]
    ht = len({i for i in top_t if t[i] > 0} & set(rel))     # a score of 0 is not a match
    he = len({i for i in top_e if e[i] > 0} & set(rel))
    k = min(3, len(rel))
    hits['TF-IDF'] += ht
    hits['embedding'] += he
    total += k
    print(f"{q:32s} TF-IDF top 3 {' '.join(f'#{i}:{t[i]:.2f}' for i in top_t)} ({ht}/{k})   "
          f"embedding top 3 {' '.join(f'#{i}:{e[i]:.2f}' for i in top_e)} ({he}/{k})")
print(f"relevant sentences found in the top 3: TF-IDF {hits['TF-IDF']}/{total}, embedding {hits['embedding']}/{total}")

b64 = lambda q: base64.b64encode(q.astype(np.int8).tobytes()).decode()
data = {
    'model': MODEL, 'dim': int(E.shape[1]), 'sentences': SENTENCES,
    'queries': [{'q': q, 'rel': rel} for q, rel in QUERIES],
    'E': b64(qE), 'Es': np.round(sE, 7).tolist(), 'Q': b64(qQ), 'Qs': np.round(sQ, 7).tolist(),
}
js = Path(__file__).resolve().parents[5] / 'docs/assets/javascripts/sims/sim-sentence-search.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}")