"""Check that the Python reference (roommap/) and the browser maths (web/embed.js, web/mapmath.js) agree.

    python tests/parity.py              # needs node on PATH; exits non-zero on any mismatch

Runs both pipelines on the 60 sample phrases (and a second, smaller list) with the fallback embedder and
compares, in order of how exact the two can be:

* embeddings: identical to 1e-12 (same hashes, same weights, same arithmetic)
* PCA: all pairwise distances between 10-d scores within 1e-7 (same subspace; signs may differ)
* mulberry32: the same first three random numbers for seed 42
* k, labels (up to relabelling), silhouette curve and cluster names: identical
* t-SNE: the first 10 iterations identical to 1e-9 relative (after 30 the gap is already ~1e-5). The final layouts are NOT expected to match
  point for point: t-SNE's gradient descent (momentum, early exaggeration, sign-switched gains) is chaotic,
  and a 1e-12 nudge of the input in Python alone changes the final picture (see the how-it-works page).
  numpy sums pairwise and JavaScript sums sequentially, which is nudge enough. What is checked instead is
  that both reach layouts of the same quality on the 60-phrase sample: KL divergence within 0.15, and
  trustworthiness and neighbourhood preservation (k = 5) within 0.1.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import tempfile

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from roommap import cluster, embed, metrics, pca, pipeline, tsne  # noqa: E402
from roommap.sample import SAMPLE  # noqa: E402

SMALL = ["Sourdough bread", "baking croissants", "jazz piano", "learning guitar", "neural networks", "diffusion models",
         "trail running", "a marathon in spring", "chess puzzles", "board games", "pottery", "knitting a scarf"]


class Mismatch(Exception):
    pass


def check(name, ok, detail=""):
    print(("ok   " if ok else "FAIL ") + name + (("  " + str(detail)) if detail else ""))
    if not ok:
        raise Mismatch(name)


def pairwise(X):
    X = np.asarray(X)
    return np.sqrt(np.maximum(((X[:, None, :] - X[None, :, :]) ** 2).sum(-1), 0))


def same_partition(a, b):
    return len({(x, y) for x, y in zip(a, b)}) == len(set(a)) == len(set(b))


def run_js(texts):
    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as f:
        f.write("\n".join(texts))
    try:
        out = subprocess.run(["node", os.path.join(ROOT, "tests", "parity.js"), f.name], capture_output=True, text=True, check=True).stdout
    finally:
        os.unlink(f.name)
    return json.loads(out)


def compare(texts, label):
    print(f"--- {label}: {len(texts)} phrases")
    js = run_js(texts)
    py = pipeline.to_json(pipeline.run(texts))
    V_js, V_py = np.array(js["vectors"]), np.array(py["vectors"])
    check("embeddings identical", V_js.shape == V_py.shape and np.abs(V_js - V_py).max() < 1e-12, f"max diff {np.abs(V_js - V_py).max():.2e}")
    d = np.abs(pairwise(js["pcs"]) - pairwise(py["pcs"])).max()
    check("PCA distances", d < 1e-7, f"max diff {d:.2e}")
    check("explained variance", np.abs(np.array(js["explained"]) - np.array(py["explained"])).max() < 1e-9)
    rnd = cluster.mulberry32(42)
    check("mulberry32", all(abs(rnd() - v) < 1e-15 for v in js["mulberry"]), js["mulberry"])
    check("k", js["k"] == py["k"], (js["k"], py["k"]))
    check("labels (up to relabelling)", same_partition(js["labels"], py["labels"]))
    sj, sp = js["silhouettes"], py["silhouettes"]
    check("silhouette curve", len(sj) == len(sp) and all(a[0] == b[0] and abs(a[1] - b[1]) < 1e-9 for a, b in zip(sj, sp)))
    check("cluster names", sorted(js["names"]) == sorted(py["names"]), (js["names"], py["names"]))
    # early t-SNE iterations, recomputed here from the Python PCA scores (sign-aligned to the JS ones)
    pcs_py, pcs_js = np.array(py["pcs"]), np.array(js["pcs"])
    if len(texts) >= 5:
        signs = np.sign((pcs_py * pcs_js).sum(0)); signs[signs == 0] = 1
        pcs_al = pcs_py * signs
        P = tsne.affinities(tsne.squared_distances(pcs_al), tsne.default_perplexity(len(texts)))
        t = tsne.TSNE(P, tsne.initial_positions(pcs_al)); t.step(10)
        Y30 = np.array(js_early(texts, 10))
        rel = np.abs(t.Y - Y30).max() / max(np.abs(t.Y).max(), 1e-300)
        check("t-SNE first 10 iterations", rel < 1e-9, f"relative diff {rel:.2e}")
    if len(texts) >= 30:
        check("KL divergence (quality, not identity)", abs(js["kl"] - py["kl"]) < 0.15, (round(js["kl"], 4), round(py["kl"], 4)))
        check("trustworthiness (quality)", abs(js["trustworthiness"] - py["trustworthiness"]) < 0.1, (round(js["trustworthiness"], 4), round(py["trustworthiness"], 4)))
        check("neighbourhood preservation (quality)", abs(js["preservation"] - py["preservation"]) < 0.1, (round(js["preservation"], 4), round(py["preservation"], 4)))
        Yj, Yp = np.array(js["layout"]), np.array(py["layout"])
        nj, npy = metrics.knn(metrics.pairwise_sq(Yj), 5), metrics.knn(metrics.pairwise_sq(Yp), 5)
        share = np.mean([len(set(nj[i]) & set(npy[i])) / 5 for i in range(len(texts))])
        print(f"info map neighbours shared between the two final layouts: {share:.2f} (not a check; see docstring)")
    return js, py


def js_early(texts, iters=30):
    """Positions after `iters` t-SNE iterations from the JavaScript."""
    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as f:
        f.write("\n".join(texts))
    try:
        out = subprocess.run(["node", os.path.join(ROOT, "tests", "parity.js"), f.name, "--early", str(iters)], capture_output=True, text=True, check=True).stdout
    finally:
        os.unlink(f.name)
    return json.loads(out)


if __name__ == "__main__":
    try:
        compare(SMALL, "small list")
        js, py = compare(SAMPLE, "sample room")
        print(f"timings (ms)  python: {{{', '.join(f'{k}: {v:.0f}' for k, v in py['timing_ms'].items())}}}"
              f"   node: {{{', '.join(f'{k}: {v:.0f}' for k, v in js['timing_ms'].items())}}}")
        print("all parity checks passed")
    except Mismatch as e:
        print(f"parity FAILED: {e}")
        sys.exit(1)