"""Grouping the room: k-means with k chosen by silhouette, then a name for each group.

k-means runs in the 10-dimensional PCA space (not on the 2-D picture, which t-SNE has distorted). The
number of groups is the k in 2..8 with the best mean silhouette, where a point's silhouette compares how
close it is to its own group with how close it is to the nearest other group (+1 good, 0 on the border,
negative probably misfiled).

Names are the two words that are most over-represented in a group compared with the room as a whole:
score(word) = p_group(word) * log(p_group(word) / p_room(word)), the group's share of the KL divergence
between "phrases in this group" and "phrases in the room", with function words excluded.

Randomness is a seeded mulberry32 generator (32-bit integer arithmetic), the same in Python and JavaScript,
so k-means++ picks the same starting centres in both.
"""
from __future__ import annotations

import math

import numpy as np

from .tsne import squared_distances
from .embed import MARKER, STOPWORDS, TOPIC_NAMES, VOCAB, WORD_TOPICS, tokenize

MASK = 0xFFFFFFFF


def mulberry32(seed: int):
    """The same tiny PRNG as web/map.js: returns a function giving floats in [0, 1)."""
    a = seed & MASK

    def imul(x: int, y: int) -> int:
        return (x * y) & MASK

    def rnd() -> float:
        nonlocal a
        a = (a + 0x6D2B79F5) & MASK
        t = a
        t = imul(t ^ (t >> 15), t | 1)
        t ^= (t + imul(t ^ (t >> 7), t | 61)) & MASK
        t &= MASK
        return ((t ^ (t >> 14)) & MASK) / 4294967296.0

    return rnd


def kmeans_pp_init(X: np.ndarray, k: int, rnd) -> np.ndarray:
    n = X.shape[0]
    centres = [int(math.floor(rnd() * n))]
    d2 = np.sum((X - X[centres[0]]) ** 2, axis=1)
    for _ in range(1, k):
        total = float(d2.sum())
        if total <= 0:
            centres.append(int(math.floor(rnd() * n)))
        else:
            r = rnd() * total
            acc = 0.0
            pick = n - 1
            for i in range(n):
                acc += float(d2[i])
                if acc >= r:
                    pick = i
                    break
            centres.append(pick)
        d2 = np.minimum(d2, np.sum((X - X[centres[-1]]) ** 2, axis=1))
    return X[centres].copy()


def kmeans(X: np.ndarray, k: int, seed: int = 1, restarts: int = 3, max_iter: int = 100) -> tuple[np.ndarray, float]:
    """Lloyd's algorithm from k-means++ starts; keeps the best of `restarts` runs. Returns (labels, inertia)."""
    n = X.shape[0]
    best_labels, best_inertia = None, math.inf
    for r in range(restarts):
        rnd = mulberry32(seed * 1000 + k * 10 + r)
        C = kmeans_pp_init(X, k, rnd)
        labels = np.full(n, -1)
        for _ in range(max_iter):
            D = np.sum(X * X, axis=1)[:, None] + np.sum(C * C, axis=1)[None, :] - 2.0 * (X @ C.T)
            new = np.argmin(D, axis=1)
            if np.array_equal(new, labels):
                break
            labels = new
            for j in range(k):
                m = labels == j
                if m.any():
                    C[j] = X[m].mean(axis=0)
        inertia = float(np.sum((X - C[labels]) ** 2))
        if inertia < best_inertia - 1e-12:
            best_inertia, best_labels = inertia, labels.copy()
    return best_labels, best_inertia


def silhouette(X: np.ndarray, labels: np.ndarray) -> float:
    """Mean silhouette over all points (Euclidean). Points alone in their group score 0."""
    n = X.shape[0]
    D = np.sqrt(squared_distances(X))
    ks = np.unique(labels)
    s = np.zeros(n)
    for i in range(n):
        own = labels == labels[i]
        if own.sum() <= 1:
            s[i] = 0.0
            continue
        a = D[i, own].sum() / (own.sum() - 1)
        b = math.inf
        for kk in ks:
            if kk == labels[i]:
                continue
            m = labels == kk
            b = min(b, D[i, m].mean())
        s[i] = (b - a) / max(a, b) if max(a, b) > 0 else 0.0
    return float(s.mean())


def choose_k(X: np.ndarray, kmin: int = 2, kmax: int = 8, seed: int = 1) -> tuple[int, np.ndarray, list[tuple[int, float]]]:
    """Try k = kmin..kmax (capped at n // 4), pick the best silhouette. Returns (k, labels, [(k, score)...]).
    Rooms with fewer than 8 phrases are one group."""
    n = X.shape[0]
    kmax = min(kmax, n // 4)
    if n < 8 or kmax < kmin:
        return 1, np.zeros(n, dtype=int), []
    best_k, best_labels, best_s, scores = 1, np.zeros(n, dtype=int), -math.inf, []
    for k in range(kmin, kmax + 1):
        labels, _ = kmeans(X, k, seed=seed)
        if len(np.unique(labels)) < k:
            continue
        s = silhouette(X, labels)
        scores.append((k, s))
        if s > best_s + 1e-12:
            best_k, best_labels, best_s = k, labels, s
    return best_k, best_labels, scores


def term_kind(w: str) -> int:
    """Tie-break order for names: topic tags, then topical words, then unknown words, then everyday words."""
    if w.startswith("#"):
        return 0
    r = VOCAB.get(w)
    if r is None:
        return 2
    return 1 if r >= MARKER else 3


def doc_terms(text: str) -> set[str]:
    """The terms a phrase is named by: its content words plus a "#topic" tag for every topic it touches."""
    words = [t for t in tokenize(text) if t not in STOPWORDS and len(t) > 1 and not t.isdigit()]
    terms = set(words)
    for w in words:
        for t in WORD_TOPICS.get(w, ()):
            terms.add("#" + TOPIC_NAMES[t])
    return terms


def name_clusters(texts: list[str], labels: np.ndarray, top: int = 2) -> list[str]:
    """One short name per group: its `top` most distinctive terms, joined with a middle dot."""
    n = len(texts)
    k = int(labels.max()) + 1 if n else 0
    docs = [doc_terms(s) for s in texts]
    room: dict[str, int] = {}
    for d in docs:
        for w in d:
            room[w] = room.get(w, 0) + 1
    names = []
    for c in range(k):
        members = [docs[i] for i in range(n) if labels[i] == c]
        if not members:
            names.append("")
            continue
        counts: dict[str, int] = {}
        for d in members:
            for w in d:
                counts[w] = counts.get(w, 0) + 1
        scored = []
        for w, cnt in counts.items():
            pc = cnt / len(members)
            pr = room[w] / n
            scored.append((-(pc * math.log(pc / pr)), -cnt, term_kind(w), w))
        scored.sort()
        picked: list[str] = []
        for _, _, _, w in scored:
            w = w.lstrip("#")
            if w not in picked:
                picked.append(w)
            if len(picked) == top:
                break
        names.append(" · ".join(picked))
    return names