"""The routing policies and the threshold search.

A *cascade* is a list of thresholds, one per model in cheap-to-expensive order. A request is sent
to the first model in the list that is not skipped; if that model's confidence is at least its
threshold the answer is kept, otherwise the request moves on to the next model and the cost of
both calls is paid. The last model that is not skipped always answers. `None` means skip;
0 means accept everything (so later models are never reached).

A *router* picks one model up front from a prediction of which models will be right, so it can
only use what is known before the call: the prompt text.

Everything is deterministic, and web/router.js is a line-for-line port. Where the two could
differ (random numbers, rounding, ties) the code is written so that both come out identical.
"""
from __future__ import annotations

import math
from typing import Optional

import numpy as np

from .data import EvalLog

SKIP = None


# ------------------------------------------------------------------ random split --

def mulberry32(seed: int):
    """The same tiny PRNG as web/router.js, so both pick the same calibration rows."""
    a = seed & 0xFFFFFFFF

    def rnd():
        nonlocal a
        a = (a + 0x6D2B79F5) & 0xFFFFFFFF
        t = a
        t = ((t ^ (t >> 15)) * (t | 1)) & 0xFFFFFFFF
        t = (((t + (((t ^ (t >> 7)) * (t | 61)) & 0xFFFFFFFF)) & 0xFFFFFFFF) ^ t) & 0xFFFFFFFF
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296.0

    return rnd


def split_indices(n: int, cal_frac: float, seed: int):
    """Shuffle 0..n-1 (Fisher-Yates) and cut: first part calibrates, the rest certifies."""
    rnd = mulberry32(seed)
    perm = list(range(n))
    for i in range(n - 1, 0, -1):
        j = int(math.floor(rnd() * (i + 1)))
        perm[i], perm[j] = perm[j], perm[i]
    k = int(math.floor(n * cal_frac + 0.5))
    k = min(max(k, 1), n - 1)
    return np.array(perm[:k], dtype=int), np.array(perm[k:], dtype=int)


# ---------------------------------------------------------------- evaluation --

def _mean_lat(log: EvalLog, idx, v):
    return float(v.sum() / len(idx)) if log.latency is not None else None


def evaluate_single(log: EvalLog, idx: np.ndarray, i: int) -> dict:
    n = len(idx)
    return {
        "quality": float(log.correct[idx, i].sum() / n),
        "qualitySum": float(log.correct[idx, i].sum()),
        "cost": float(log.cost[idx, i].sum() / n),
        "latency": float(log.latency[idx, i].sum() / n) if log.latency is not None else None,
        "shares": [1.0 if j == i else 0.0 for j in range(log.m)],
        "calls": [1.0 if j == i else 0.0 for j in range(log.m)],
        "maxCost": float(log.cost[idx, i].max()),
    }


def canonical(thr):
    """Last active model always accepts: write its threshold as 0 so equal cascades look equal."""
    thr = list(thr)
    last = max((i for i, t in enumerate(thr) if t is not SKIP), default=-1)
    if last >= 0:
        thr[last] = 0.0
    return thr


def evaluate_cascade(log: EvalLog, idx: np.ndarray, thr, conf: np.ndarray) -> dict:
    """Run the cascade on the rows idx. conf is the (n, m) gating signal (reported or learned)."""
    n, m = len(idx), log.m
    active = [i for i, t in enumerate(thr) if t is not SKIP]
    if not active:
        raise ValueError("a cascade needs at least one model")
    remaining = np.ones(n, dtype=bool)
    answered = np.full(n, -1, dtype=int)
    cost = np.zeros(n)
    lat = np.zeros(n)
    calls = [0.0] * m
    for k, i in enumerate(active):
        if not remaining.any():
            break
        rows = idx[remaining]
        cost[remaining] += log.cost[rows, i]
        if log.latency is not None:
            lat[remaining] += log.latency[rows, i]
        calls[i] = float(remaining.sum() / n)
        if k == len(active) - 1:
            acc = remaining.copy()
        else:
            acc = remaining & (conf[idx, i] >= thr[i])
        answered[acc] = i
        remaining &= ~acc
    q = log.correct[idx, answered]
    shares = [float((answered == i).sum() / n) for i in range(m)]
    return {
        "quality": float(q.sum() / n),
        "qualitySum": float(q.sum()),
        "cost": float(cost.sum() / n),
        "latency": float(lat.sum() / n) if log.latency is not None else None,
        "shares": shares,
        "calls": calls,
        "maxCost": float(log.cost[idx][:, active].sum(axis=1).max()),
        "thresholds": canonical(thr),
    }


def evaluate_router(log: EvalLog, idx: np.ndarray, tau: float, pred: np.ndarray) -> dict:
    """Send each prompt to the cheapest model whose predicted P(correct) >= tau, else the last model."""
    n, m = len(idx), log.m
    p = pred[idx]
    ok = p >= tau
    choice = np.where(ok.any(axis=1), ok.argmax(axis=1), m - 1)
    q = log.correct[idx, choice]
    c = log.cost[idx, choice]
    return {
        "quality": float(q.sum() / n),
        "qualitySum": float(q.sum()),
        "cost": float(c.sum() / n),
        "latency": float(log.latency[idx, choice].sum() / n) if log.latency is not None else None,
        "shares": [float((choice == i).sum() / n) for i in range(m)],
        "calls": [float((choice == i).sum() / n) for i in range(m)],
        "maxCost": float(log.cost[idx].max()),
        "tau": float(tau),
    }


# -------------------------------------------------------------------- oracle --

def oracle(log: EvalLog, idx: np.ndarray) -> dict:
    """Per prompt, the cheapest model that is right (or the cheapest model if none is).

    The curve: start from 'cheapest model for everything' and buy the prompts' corrections in
    order of price, so it is the lowest cost any policy can reach at each quality level.
    """
    n, m = len(idx), log.m
    cor = log.correct[idx]
    cost = log.cost[idx]
    best = np.full(n, -1, dtype=int)
    for i in range(m - 1, -1, -1):          # cheapest correct wins (scan expensive -> cheap)
        best = np.where(cor[:, i] >= 1.0, i, best)
    base_cost = cost[:, 0]
    base_q = cor[:, 0]
    choice = np.where(best >= 0, best, 0)
    full_cost = cost[np.arange(n), choice]
    full_q = cor[np.arange(n), choice]
    # the price of each correction, and its gain in quality
    price = full_cost - base_cost
    gain = full_q - base_q
    scale = float(log.cost.mean(axis=0).max())
    order = sorted(range(n), key=lambda j: (_key(float(price[j]), scale), -_key(float(gain[j]), 1.0), j))
    pts = [{"cost": float(base_cost.sum() / n), "quality": float(base_q.sum() / n)}]
    cs, qs = float(base_cost.sum()), float(base_q.sum())
    for j in order:
        if gain[j] <= 0:
            continue
        cs += float(price[j])
        qs += float(gain[j])
        pts.append({"cost": cs / n, "quality": qs / n})
    shares = [float((choice == i).sum() / n) for i in range(m)]
    return {
        "quality": float(full_q.sum() / n),
        "qualitySum": float(full_q.sum()),
        "cost": float(full_cost.sum() / n),
        "latency": float(log.latency[idx, choice].sum() / n) if log.latency is not None else None,
        "shares": shares,
        "calls": shares,
        "maxCost": float(full_cost.max()),
        "curve": pts,
        "noneRight": float((best < 0).sum() / n),
    }


def pareto(points):
    """Indices of the (cost, quality) points that no other point beats on both."""
    keep = []
    for i, p in enumerate(points):
        dominated = any((q["cost"] <= p["cost"] and q["quality"] >= p["quality"] and
                         (q["cost"] < p["cost"] or q["quality"] > p["quality"])) for j, q in enumerate(points) if j != i)
        if not dominated:
            keep.append(i)
    return keep


# ------------------------------------------------------------ threshold grids --

def quantile_grid(values: np.ndarray, g: int):
    """0 (accept all) plus g-1 quantiles of the observed confidences, duplicates removed."""
    s = np.sort(values)
    n = len(s)
    out = [0.0]
    for k in range(1, g):
        v = float(s[int(math.floor((n - 1) * k / (g - 1)))])
        if v > out[-1]:
            out.append(v)
    return out


def grid_size(m: int, budget: int = 20000, g_max: int = 24) -> int:
    """Largest grid such that the exhaustive enumeration stays under `budget` cascades."""
    g = g_max
    while g > 3:
        total = sum((g + 1) ** L for L in range(m))
        if total <= budget:
            break
        g -= 1
    return g


# ----------------------------------------------------- exhaustive enumeration --

def _key(v: float, scale: float) -> int:
    """Costs and qualities are compared after rounding, so numpy's pairwise sums and JS's plain
    loops cannot disagree about a tie."""
    return int(math.floor(v / scale * 1e9 + 0.5))


def enumerate_cascades(log: EvalLog, idx: np.ndarray, conf: np.ndarray, grids):
    """Every cascade on the grids, with its calibration cost and quality.

    Returns (thresholds list, cost per prompt array, quality array). A model with no accepted
    rows (useless call) or with every remaining row accepted (then it is the last model) is not
    enumerated twice.
    """
    n, m = len(idx), log.m
    conf_c = conf[idx]
    cost_c = log.cost[idx]
    cor_c = log.correct[idx]
    thrs, costs, quals = [], [], []

    def leaf(thr, csum, qsum):
        thrs.append(thr)
        costs.append(csum / n)
        quals.append(qsum / n)

    def rec(i, remaining, thr, csum, qsum):
        rows_cost = float(cost_c[remaining, i].sum())
        # model i is the last active one: it answers everything that is left
        leaf(thr + [0.0] + [SKIP] * (m - i - 1), csum + rows_cost, qsum + float(cor_c[remaining, i].sum()))
        if i == m - 1:
            return
        # skip model i
        rec(i + 1, remaining, thr + [SKIP], csum, qsum)
        # or gate on a threshold
        for t in grids[i][1:]:
            acc = remaining & (conf_c[:, i] >= t)
            n_acc = int(acc.sum())
            if n_acc == 0 or n_acc == int(remaining.sum()):
                continue
            rec(i + 1, remaining & ~acc, thr + [t], csum + rows_cost, qsum + float(cor_c[acc, i].sum()))

    rec(0, np.ones(n, dtype=bool), [], 0.0, 0.0)
    return thrs, np.array(costs), np.array(quals)


def pick_best(costs, quals, target: float, scale: float):
    """Cheapest cascade with quality >= target; ties go to the higher quality, then the first
    enumerated. With no feasible cascade: the highest quality, then the cheapest."""
    ck = np.floor(np.asarray(costs, dtype=float) / scale * 1e9 + 0.5).astype(np.int64)
    qk = np.floor(np.asarray(quals, dtype=float) * 1e9 + 0.5).astype(np.int64)
    tk = _key(target, 1.0)
    feas = qk >= tk
    idx = np.arange(len(ck))
    if feas.any():
        cand = idx[feas]
        order = np.lexsort((cand, -qk[cand], ck[cand]))   # primary key last: cost, then -quality, then index
        return int(cand[order[0]]), True
    order = np.lexsort((idx, ck, -qk))
    return int(order[0]), False


def refine(log: EvalLog, idx: np.ndarray, conf: np.ndarray, thr, target: float, fine_grids, scale: float, passes: int = 2):
    """Coordinate descent on a finer grid, starting from the enumerated optimum."""
    thr = list(thr)
    cur = evaluate_cascade(log, idx, thr, conf)
    tk = _key(target, 1.0)

    def better(a, b):  # is a better than b for this target?
        fa, fb = _key(a["quality"], 1.0) >= tk, _key(b["quality"], 1.0) >= tk
        if fa != fb:
            return fa
        if fa:
            return (_key(a["cost"], scale), -_key(a["quality"], 1.0)) < (_key(b["cost"], scale), -_key(b["quality"], 1.0))
        return (-_key(a["quality"], 1.0), _key(a["cost"], scale)) < (-_key(b["quality"], 1.0), _key(b["cost"], scale))

    for _ in range(passes):
        changed = False
        for i in range(log.m):
            for t in [SKIP] + list(fine_grids[i]):
                cand = thr[:i] + [t] + thr[i + 1:]
                if all(x is SKIP for x in cand):
                    continue
                r = evaluate_cascade(log, idx, cand, conf)
                if better(r, cur):
                    cur, thr, changed = r, canonical(cand), True
        if not changed:
            break
    return thr, cur


# ------------------------------------------------------------ learned signal --

FEATURE_DIM = 256


def fnv1a(s: str) -> int:
    h = 0x811C9DC5
    for ch in s.encode("utf-8"):
        h ^= ch
        h = (h * 0x01000193) & 0xFFFFFFFF
    return h


def tokens(text: str):
    out, cur = [], []
    for ch in text.lower():
        if ch.isalnum() and ord(ch) < 128:
            cur.append(ch)
        else:
            if cur:
                out.append("".join(cur))
            cur = []
    if cur:
        out.append("".join(cur))
    return out


def features(prompts) -> np.ndarray:
    """Bias, length cues and a hashed bag of words (FNV-1a into 256 buckets, 0/1)."""
    n = len(prompts)
    X = np.zeros((n, 5 + FEATURE_DIM))
    for r, p in enumerate(prompts):
        p = p or ""
        toks = tokens(p)
        X[r, 0] = 1.0
        X[r, 1] = math.log(1.0 + len(p)) / 5.0
        X[r, 2] = math.log(1.0 + len(toks)) / 5.0
        X[r, 3] = sum(1 for ch in p if "0" <= ch <= "9") / (len(p) + 1.0)
        X[r, 4] = 1.0 if "?" in p else 0.0
        for t in toks:
            X[r, 5 + fnv1a(t) % FEATURE_DIM] = 1.0
    return X


def fit_logistic(X: np.ndarray, y: np.ndarray, iters: int = 300, lr: float = 0.5, l2: float = 1e-3) -> np.ndarray:
    """Full-batch gradient descent on L2-regularised logistic loss (soft labels allowed)."""
    n, d = X.shape
    w = np.zeros(d)
    for _ in range(iters):
        p = 1.0 / (1.0 + np.exp(-(X @ w)))
        g = X.T @ (p - y) / n
        g[1:] += l2 * w[1:]
        w -= lr * g
    return w


def predict(X: np.ndarray, w: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-(X @ w)))


def learn_confidence(log: EvalLog, cal_idx: np.ndarray, folds: int = 5) -> np.ndarray:
    """P(model i is right | prompt text), as an (n, m) matrix.

    Calibration rows get out-of-fold predictions (fold = position in the shuffled calibration
    list mod `folds`), so the thresholds are tuned on predictions the router did not fit on.
    Every other row (the held-out split) gets the prediction of a model fit on all calibration rows.
    """
    if log.prompts is None:
        raise ValueError("learning a router needs a prompt text column")
    X = features(log.prompts)
    n, m = log.n, log.m
    pred = np.zeros((n, m))
    cal_set = np.zeros(n, dtype=bool)
    cal_set[cal_idx] = True
    for i in range(m):
        y = log.correct[:, i]
        w_all = fit_logistic(X[cal_idx], y[cal_idx])
        pred[~cal_set, i] = predict(X[~cal_set], w_all)
        for f in range(folds):
            pos = np.arange(len(cal_idx))
            test = cal_idx[pos % folds == f]
            train = cal_idx[pos % folds != f]
            if len(test) == 0 or len(train) == 0:
                continue
            w = fit_logistic(X[train], y[train])
            pred[test, i] = predict(X[test], w)
    return pred