"""
Glue between the demo page and the real OVec library.

In the browser this file runs inside Pyodide next to an unmodified copy of
``vector_db/``. It runs the same way under plain CPython, which is how
``demo/test_bridge.py`` checks it.

Nothing here reimplements search. Every result comes from ``VectorDB`` itself.
The figures the page reports *about* a search (how many vectors were scored,
which clusters were probed, which branch of ``search()`` ran) are taken by
watching the library's own calls to ``compute_distances`` and
``get_top_k_indices`` while it runs. The originals are put back as soon as the
search returns, and timings are taken with nothing patched.
"""

from contextlib import contextmanager
from time import perf_counter

import numpy as np

import vector_db.core as _core
import vector_db.index as _index
from vector_db import VectorDB
from vector_db.core import validate_vector

DIM = 2
GRID = 256            # resolution of the cluster-region map, per side
EXTENT = 1.0          # the plane shown is [-EXTENT, EXTENT]^2
TIME_BUDGET_S = 0.012  # per method, per query: repeat until this much time has passed
MAX_TIMED_RUNS = 400


# --------------------------------------------------------------------------
# Observing the library
# --------------------------------------------------------------------------

class _Observation:
    def __init__(self):
        self.brute_force_scored = 0  # vectors scored by core.compute_distances
        self.ivf_scored = 0          # candidate vectors scored inside IVFIndex.search
        self.centroids_ranked = 0    # centroids scored to choose the clusters to probe
        self.probed = []             # cluster ids, nearest first
        self._centroid_scores = None

    @property
    def path(self):
        if self.ivf_scored or self.centroids_ranked:
            return "ivf"
        if self.brute_force_scored:
            return "brute_force"
        return "none"


@contextmanager
def _observe(index=None):
    """Count what the library computes during the ``with`` block."""
    obs = _Observation()
    centroids = index.kmeans.centroids if index is not None and index.kmeans is not None else None
    core_cd = _core.compute_distances
    index_cd = _index.compute_distances
    index_topk = _index.get_top_k_indices

    def watched_core_cd(query, vectors, metric="cosine"):
        obs.brute_force_scored += len(vectors)
        return core_cd(query, vectors, metric)

    def watched_index_cd(query, vectors, metric="cosine"):
        scores = index_cd(query, vectors, metric)
        if centroids is not None and vectors is centroids:
            obs.centroids_ranked += len(vectors)
            obs._centroid_scores = scores
        else:
            obs.ivf_scored += len(vectors)
        return scores

    def watched_index_topk(scores, k):
        top = index_topk(scores, k)
        if scores is obs._centroid_scores:
            obs.probed = [int(c) for c in top]
        return top

    _core.compute_distances = watched_core_cd
    _index.compute_distances = watched_index_cd
    _index.get_top_k_indices = watched_index_topk
    try:
        yield obs
    finally:
        _core.compute_distances = core_cd
        _index.compute_distances = index_cd
        _index.get_top_k_indices = index_topk


def _mean_ms(fn):
    """Mean wall time of ``fn()`` in ms, repeated until the time budget is spent.

    Browsers coarsen timers (often to 0.1 ms), so a single microsecond-scale
    search cannot be timed on its own.
    """
    runs = 0
    start = perf_counter()
    while True:
        fn()
        runs += 1
        elapsed = perf_counter() - start
        if elapsed >= TIME_BUDGET_S or runs >= MAX_TIMED_RUNS:
            return elapsed * 1000.0 / runs, runs


# --------------------------------------------------------------------------
# Data
# --------------------------------------------------------------------------

def make_points(n, seed):
    """A few thousand 2-D points: anisotropic Gaussian blobs plus some background."""
    rng = np.random.default_rng(seed)
    n_blobs = int(rng.integers(6, 10))
    centers = rng.uniform(-0.62, 0.62, (n_blobs, DIM))
    weights = rng.dirichlet(np.full(n_blobs, 3.0))
    n_noise = int(n * 0.04)
    counts = rng.multinomial(n - n_noise, weights)

    chunks = []
    for center, count in zip(centers, counts):
        angle = rng.uniform(0, np.pi)
        rot = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
        axes = rng.uniform(0.07, 0.22, DIM)
        chunks.append(center + (rng.normal(size=(count, DIM)) * axes) @ rot.T)
    chunks.append(rng.uniform(-0.95, 0.95, (n_noise, DIM)))
    points = np.vstack(chunks)

    # Keep everything inside the visible plane by redrawing strays from the background.
    outside = np.any(np.abs(points) > 0.97, axis=1)
    points[outside] = rng.uniform(-0.95, 0.95, (int(outside.sum()), DIM))
    rng.shuffle(points)
    return points.astype(np.float32)


# --------------------------------------------------------------------------
# Session state (one per page)
# --------------------------------------------------------------------------

class _Session:
    def __init__(self):
        self.points = None
        self.db = None
        self.metric = "euclidean"
        self.nlist = 32
        self.nprobe = 2


S = _Session()


def _new_db():
    """A fresh VectorDB holding the current points. No index: search() is exact."""
    db = VectorDB(dimension=DIM, metric=S.metric, index_type="ivf",
                  n_clusters=S.nlist, nprobe=S.nprobe)
    db.batch_insert([str(i) for i in range(len(S.points))], S.points)
    S.db = db


def _indexed():
    return S.db is not None and S.db._index_built and S.db._index is not None


def _state():
    return {
        "n": len(S.db),
        "metric": S.metric,
        "indexed": _indexed(),
        "nlist": S.nlist,
        "nprobe": S.nprobe,
    }


def load_data(n, seed, metric):
    """Generate points and put them in a new VectorDB."""
    S.points = make_points(int(n), int(seed))
    S.metric = str(metric)
    _new_db()
    return {"points": S.points.tobytes(), **_state()}


def set_metric(metric):
    """The metric is a VectorDB constructor argument, so this means a new database."""
    S.metric = str(metric)
    _new_db()
    return _state()


def build_index(nlist, nprobe, seed):
    """Call db.build_index(). k-means draws from NumPy's global RNG, so seed it."""
    S.nlist = int(nlist)
    S.nprobe = max(1, int(nprobe))
    S.db.index_params["n_clusters"] = S.nlist
    S.db.index_params["nprobe"] = S.nprobe
    np.random.seed(int(seed))
    start = perf_counter()
    S.db.build_index()
    build_ms = (perf_counter() - start) * 1000.0

    index = S.db._index
    labels = np.zeros(len(S.points), dtype=np.int16)
    for cluster, members in index.inverted_lists.items():
        labels[members] = cluster

    # The region each cluster owns, from the index's own KMeans.predict.
    ticks = (np.arange(GRID, dtype=np.float32) + 0.5) * (2 * EXTENT / GRID) - EXTENT
    gx, gy = np.meshgrid(ticks, ticks[::-1])  # row 0 is the top of the plane
    grid = np.column_stack([gx.ravel(), gy.ravel()]).astype(np.float32)
    regions = np.concatenate([index.kmeans.predict(grid[i:i + 4096])
                              for i in range(0, len(grid), 4096)]).astype(np.int16)

    nlist_effective = int(index.kmeans.n_clusters)
    if S.nprobe > nlist_effective:
        set_nprobe(nlist_effective)
    return {
        **_state(),
        "build_ms": build_ms,
        "nlist_effective": nlist_effective,
        "centroids": index.kmeans.centroids.astype(np.float32).tobytes(),
        "labels": labels.tobytes(),
        "regions": regions.tobytes(),
        "grid": GRID,
        "sizes": [len(index.inverted_lists[c]) for c in range(nlist_effective)],
    }


def set_nprobe(nprobe):
    """nprobe is read on every IVFIndex.search call, so it can change without a rebuild."""
    S.nprobe = max(1, int(nprobe))
    S.db.index_params["nprobe"] = S.nprobe
    if _indexed():
        S.db._index.nprobe = S.nprobe
    return _state()


def _ids(results):
    return [int(r["id"]) for r in results]


def query(x, y, k):
    """Search the point (x, y) the way the database would, and with brute force."""
    k = int(k)
    q = np.array([x, y], dtype=np.float32)
    db = S.db

    with _observe(db._index if _indexed() else None) as seen:
        results = db.search(q, top_k=k)          # the public call; dispatch decides

    qv = validate_vector(q, DIM)
    out = {"path": seen.path, "k": k, **_state()}

    if seen.path == "brute_force":
        exact, exact_seen = results, seen
    else:
        with _observe() as exact_seen:
            exact = db._search_brute_force(qv, k, True)  # what search() runs without an index
    ms, runs = _mean_ms(lambda: db._search_brute_force(qv, k, True))
    out["exact"] = {
        "ids": _ids(exact),
        "scores": [r["score"] for r in exact],
        "scored": exact_seen.brute_force_scored,
        "ms": ms,
        "runs": runs,
    }

    if seen.path == "ivf":
        ms, runs = _mean_ms(lambda: db._search_ivf(qv, k, True))
        truth = set(out["exact"]["ids"])
        found = _ids(results)
        out["ivf"] = {
            "ids": found,
            "scores": [r["score"] for r in results],
            "scored": seen.ivf_scored,
            "centroids_ranked": seen.centroids_ranked,
            "probed": seen.probed,
            "ms": ms,
            "runs": runs,
        }
        out["hits"] = len(truth & set(found))
        out["recall"] = out["hits"] / len(truth) if truth else 1.0
    else:
        out["ivf"] = None
        out["recall"] = None
    return out


def _probe_values(nlist, current):
    values = set(range(1, min(nlist, 8) + 1))
    p = 8.0
    while p < nlist:
        p *= 1.3
        values.add(min(nlist, int(round(p))))
    values.add(nlist)
    values.add(min(nlist, max(1, current)))
    return sorted(values)


def recall_curve(k, n_queries, seed):
    """Mean recall@k and vectors scored, for a sweep of nprobe over random queries.

    Queries are drawn near the data (a random stored point plus a little noise).
    Ground truth is the database's own brute-force search.
    """
    if not _indexed():
        return None
    k = int(k)
    db, index = S.db, S.db._index
    rng = np.random.default_rng(int(seed))
    queries = S.points[rng.integers(0, len(S.points), int(n_queries))]
    queries = (queries + rng.normal(0, 0.03, queries.shape)).astype(np.float32)
    truth = [set(_ids(db._search_brute_force(q, k, False))) for q in queries]
    wanted = sum(len(t) for t in truth)

    nlist = int(index.kmeans.n_clusters)
    rows = []
    try:
        for nprobe in _probe_values(nlist, S.nprobe):
            index.nprobe = nprobe
            hits = 0
            with _observe(index) as seen:
                for q, t in zip(queries, truth):
                    hits += len(t & set(_ids(db.search(q, top_k=k, return_scores=False))))
            rows.append({
                "nprobe": nprobe,
                "recall": hits / wanted,
                "scanned": seen.ivf_scored / len(queries),
            })
    finally:
        index.nprobe = S.nprobe
    return {"k": k, "n": len(S.points), "n_queries": len(queries),
            "nlist": nlist, "metric": S.metric, "rows": rows}