"""Tests for the Python reference. Run with `python -m pytest -q` from projects/room-map/."""
import json
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest
from sklearn.decomposition import PCA as SkPCA
from sklearn.manifold import trustworthiness as sk_trust
from sklearn.metrics import silhouette_score

from roommap import cluster, embed, metrics, pca, pipeline, tsne
from roommap.sample import SAMPLE, SAMPLE_PROMPT

ROOT = Path(__file__).resolve().parent.parent


# ------------------------------------------------------------------ embed --
def test_tokenize():
    assert embed.tokenize("Sourdough. My starter's older than my car!") == ["sourdough", "my", "starter's", "older", "than", "my", "car"]
    assert embed.tokenize("  ") == []


def test_fnv1a_known_values():
    assert embed.fnv1a("") == 0x811C9DC5
    assert embed.fnv1a("a") == 0xE40C292C
    assert embed.fnv1a("foobar") == 0xBF9CF968


def test_embedding_is_unit_and_deterministic():
    v = embed.embed("Bouldering twice a week")
    assert v.shape == (embed.DIM,)
    assert abs(np.linalg.norm(v) - 1) < 1e-12
    assert np.array_equal(v, embed.embed("Bouldering twice a week"))
    assert np.linalg.norm(embed.embed("")) == 0


def test_similar_phrases_are_closer():
    a = embed.embed("baking sourdough bread every weekend")
    b = embed.embed("I bake bread and croissants")
    c = embed.embed("training neural networks on my laptop")
    assert a @ b > a @ c
    # stems share trigrams even when the word block sees different words
    assert embed.embed("climbing") @ embed.embed("climb") > embed.embed("climbing") @ embed.embed("piano")


def test_word_weights_follow_the_tiers():
    assert embed.word_weight("the") == embed.W_STOP
    assert embed.word_weight("morning") == embed.W_COMMON
    assert embed.word_weight("sourdough") == embed.W_TOPICAL
    assert embed.word_weight("zxqv") == embed.W_OOV
    assert embed.D_TOPIC == 32


# -------------------------------------------------------------------- pca --
def test_pca_matches_sklearn_up_to_sign():
    rng = np.random.default_rng(0)
    X = rng.normal(size=(40, 100)) @ rng.normal(size=(100, 100)) * np.linspace(1, 3, 100)
    ours, expl = pca.pca(X, 5)
    sk = SkPCA(5).fit(X)
    theirs = sk.transform(X)
    for j in range(5):
        assert min(np.abs(ours[:, j] - theirs[:, j]).max(), np.abs(ours[:, j] + theirs[:, j]).max()) < 1e-8
    assert np.allclose(expl, sk.explained_variance_ratio_)
    assert np.all(np.diff(expl) <= 1e-12)


def test_pca_small_inputs():
    s, e = pca.pca(np.zeros((0, 5)), 10)
    assert s.shape == (0, 0)
    s, e = pca.pca(np.array([[1.0, 2.0]]), 10)
    assert s.shape == (1, 0)
    s, e = pca.pca(np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]), 10)
    assert s.shape == (3, 2) and abs(e[0] - 1) < 1e-12


# ------------------------------------------------------------------- tsne --
def test_affinities_hit_the_perplexity():
    rng = np.random.default_rng(1)
    X = rng.normal(size=(50, 5))
    P = tsne.affinities(tsne.squared_distances(X), 10.0)
    assert abs(P.sum() - 1) < 1e-9
    assert np.allclose(P, P.T)
    # row entropy of the conditional distribution ~ log(perplexity) (check one row by rebuilding it)
    i = 3
    D = tsne.squared_distances(X)[i]
    idx = [j for j in range(50) if j != i]
    best = None
    for beta in np.geomspace(0.01, 100, 2000):
        p = np.exp(-D[idx] * beta)
        p /= p.sum()
        H = -(p[p > 0] * np.log(p[p > 0])).sum()
        if best is None or abs(H - np.log(10)) < best[0]:
            best = (abs(H - np.log(10)), beta)
    assert best[0] < 0.01


def test_layout_separates_blobs_and_reduces_kl():
    rng = np.random.default_rng(2)
    X = np.concatenate([rng.normal(size=(30, 10)) + 8 * np.eye(10)[c] for c in range(3)])
    Y, kl = tsne.layout(X)
    labels = np.repeat([0, 1, 2], 30)
    assert silhouette_score(Y, labels) > 0.5
    P = tsne.affinities(tsne.squared_distances(X), tsne.default_perplexity(90))
    t0 = tsne.TSNE(P, tsne.initial_positions(X))
    assert kl < t0.kl()


def test_layout_is_deterministic_and_centred():
    X = embed.embed_all(SAMPLE[:20])
    pcs, _ = pca.pca(X, 10)
    Y1, _ = tsne.layout(pcs, iters=100)
    Y2, _ = tsne.layout(pcs, iters=100)
    assert np.array_equal(Y1, Y2)
    assert np.abs(Y1.mean(axis=0)).max() < 1e-9


def test_tiny_rooms():
    X = embed.embed_all(SAMPLE[:3])
    Y, kl = tsne.layout(pca.pca(X, 10)[0])
    assert Y.shape == (3, 2) and kl == 0.0


# ---------------------------------------------------------------- cluster --
def test_mulberry32_known_sequence():
    # first three outputs of mulberry32(42), as JavaScript gives them (checked with node in tests/parity.py)
    rnd = cluster.mulberry32(42)
    seq = [rnd() for _ in range(3)]
    assert all(0 <= x < 1 for x in seq)
    assert cluster.mulberry32(42)() == seq[0]
    assert cluster.mulberry32(43)() != seq[0]


def test_kmeans_finds_blobs_and_silhouette_matches_sklearn():
    rng = np.random.default_rng(3)
    X = np.concatenate([rng.normal(scale=0.3, size=(25, 4)) + 3 * np.eye(4)[c] for c in range(4)])
    labels, inertia = cluster.kmeans(X, 4)
    truth = np.repeat(range(4), 25)
    # same partition up to relabelling
    assert len({(a, b) for a, b in zip(labels, truth)}) == 4
    assert abs(cluster.silhouette(X, labels) - silhouette_score(X, labels)) < 1e-9
    k, lab, scores = cluster.choose_k(X)
    assert k == 4 and max(scores, key=lambda t: t[1])[0] == 4


def test_choose_k_small_room():
    k, labels, scores = cluster.choose_k(np.random.default_rng(0).normal(size=(6, 3)))
    assert k == 1 and labels.tolist() == [0] * 6 and scores == []


def test_cluster_names_pick_distinctive_terms():
    texts = ["baking bread", "sourdough bread and croissants", "fresh bread every morning", "jazz piano", "learning guitar", "techno music"]
    labels = np.array([0, 0, 0, 1, 1, 1])
    names = cluster.name_clusters(texts, labels)
    assert names[0].split(" · ")[0] in ("food", "bread")
    assert names[1].split(" · ")[0] == "music"


# ---------------------------------------------------------------- metrics --
def test_trustworthiness_matches_sklearn():
    rng = np.random.default_rng(4)
    X = rng.normal(size=(40, 8))
    Y = X[:, :2] + 0.3 * rng.normal(size=(40, 2))
    for k in (3, 5, 7):
        assert abs(metrics.trustworthiness(X, Y, k) - sk_trust(X, Y, n_neighbors=k)) < 1e-9
    assert metrics.trustworthiness(X, X[:, :8], 5) == 1.0
    assert metrics.neighbourhood_preservation(X, X, 5) == 1.0


# --------------------------------------------------------------- pipeline --
def test_pipeline_on_the_sample():
    assert len(SAMPLE) == 60 and SAMPLE_PROMPT
    r = pipeline.run(SAMPLE)
    assert r["layout"].shape == (60, 2)
    assert r["k"] >= 2 and len(r["names"]) == r["k"]
    assert r["trustworthiness"] > 0.8
    # the cooking phrases end up together
    i, j = SAMPLE.index("Sourdough. My starter is older than my car."), SAMPLE.index("Trying to bake the perfect croissant at home")
    assert r["labels"][i] == r["labels"][j]
    assert json.dumps(pipeline.to_json(r))


def test_cli_prints_json():
    out = subprocess.run([sys.executable, "-m", "roommap.pipeline"], cwd=ROOT, capture_output=True, text=True, check=True).stdout
    assert json.loads(out)["n"] == 60