"""Cosine-similarity thresholds for embeddings: what is "similar enough"?

Why 0.2 can be a lot. For two independent directions drawn uniformly on the
unit sphere in R^d, the cosine c satisfies (1 + c)/2 ~ Beta((d-1)/2, (d-1)/2):
mean 0, standard deviation ~ 1/sqrt(d). In d = 768 that is ~0.036, so a
cosine of 0.2 is ~5.5 standard deviations above the isotropic null
(p ~ 1e-8). Intuition from 2-D/3-D ("0.2 is almost orthogonal") does not
transfer to high dimensions.

Why you still need an EMPIRICAL null. Real embedding models are anisotropic:
vectors share a common mean direction, so *unrelated* texts often have cosine
0.3-0.6. The isotropic formula then wildly overstates significance. The honest
recipe is to measure the cosine distribution on pairs you know are unrelated
(and related), optionally after mean-centring the embeddings, and pick the
threshold from that.

Picking a threshold with a precision guarantee. "Declare related if
cos >= t" is a selective-prediction problem: precision(t) = P(related |
cos >= t). We reuse the Bonferroni + Clopper-Pearson machinery from
``selective``: with probability >= 1 - delta over the labelled pairs, the
chosen t has precision >= target on new pairs drawn from the same pair
distribution. Caveats: (1) precision depends on the base rate of related
pairs; if deployment has far fewer related pairs than your labelled set,
precision drops. (2) The binomial bound assumes independent pairs; pairs
sampled from a finite labelled set share items, so treat the bound as
approximate unless pairs are disjoint (or come from independent sources).
"""

from __future__ import annotations

import numpy as np
from scipy.stats import beta

from .selective import clopper_pearson_upper


def normalize(E: np.ndarray, center: bool = False, mean: np.ndarray | None = None) -> np.ndarray:
    """Unit-normalise rows; optionally subtract a mean vector first."""
    E = np.asarray(E, dtype=float)
    if center:
        E = E - (E.mean(0) if mean is None else mean)
    return E / np.clip(np.linalg.norm(E, axis=1, keepdims=True), 1e-12, None)


def isotropic_null_pvalue(cos: float | np.ndarray, d: int) -> np.ndarray:
    """P(C >= cos) for C the cosine between two independent uniform directions in R^d."""
    a = (d - 1) / 2.0
    return beta.sf((np.asarray(cos) + 1) / 2.0, a, a)


def isotropic_null_sd(d: int) -> float:
    """Exact standard deviation of the isotropic null cosine: 1/sqrt(d)."""
    return 1.0 / np.sqrt(d)


def sample_pairs(labels: np.ndarray, n_pairs: int = 5000, seed: int = 0):
    """Sample index pairs: half same-label ("related"), half different-label."""
    rng = np.random.default_rng(seed)
    labels = np.asarray(labels)
    by = {c: np.flatnonzero(labels == c) for c in np.unique(labels)}
    multi = [c for c, ix in by.items() if len(ix) >= 2]
    pairs, rel = [], []
    half = n_pairs // 2
    while len(pairs) < half and multi:
        c = multi[rng.integers(len(multi))]
        i, j = rng.choice(by[c], 2, replace=False)
        pairs.append((i, j)); rel.append(1)
    n = len(labels)
    while len(pairs) < n_pairs:
        i, j = rng.integers(n, size=2)
        if labels[i] != labels[j]:
            pairs.append((i, j)); rel.append(0)
    return np.array(pairs), np.array(rel)


def pair_cosines(E: np.ndarray, pairs: np.ndarray) -> np.ndarray:
    En = normalize(E)
    return np.einsum("ij,ij->i", En[pairs[:, 0]], En[pairs[:, 1]])


def threshold_for_precision(
    cos: np.ndarray, related: np.ndarray, target_precision: float = 0.95,
    delta: float = 0.05, grid: np.ndarray | None = None,
) -> dict:
    """Lowest cosine threshold whose precision lower bound is >= target.

    Bound: Bonferroni over a fixed grid of thresholds + one-sided Clopper-Pearson
    on the false-positive fraction among pairs above the threshold.
    """
    if grid is None:
        grid = np.round(np.arange(-0.2, 1.0, 0.01), 4)  # fixed, data-independent
    m = len(grid)
    rows = []
    for t in grid:
        sel = cos >= t
        n_sel = int(sel.sum())
        fp = int((related[sel] == 0).sum())
        ub = clopper_pearson_upper(fp, n_sel, delta / m)
        rows.append({
            "threshold": float(t), "n_selected": n_sel,
            "precision": (n_sel - fp) / n_sel if n_sel else float("nan"),
            "precision_lower_bound": 1.0 - ub,
            "recall": float(np.sum(sel & (related == 1)) / max(1, np.sum(related == 1))),
        })
    ok = [r for r in rows if r["precision_lower_bound"] >= target_precision]
    best = min(ok, key=lambda r: r["threshold"]) if ok else None
    return {"target_precision": target_precision, "delta": delta,
            "threshold": None if best is None else best["threshold"],
            "certified": best, "table": rows}


def null_summary(cos_unrelated: np.ndarray, d: int) -> dict:
    """Compare the empirical unrelated-pair distribution to the isotropic null."""
    return {
        "dim": int(d),
        "isotropic_sd": isotropic_null_sd(d),
        "isotropic_p_at_0.2": float(isotropic_null_pvalue(0.2, d)),
        "empirical_mean": float(np.mean(cos_unrelated)),
        "empirical_sd": float(np.std(cos_unrelated)),
        "empirical_q99": float(np.quantile(cos_unrelated, 0.99)),
        "empirical_frac_above_0.2": float(np.mean(cos_unrelated >= 0.2)),
    }