"""Ranking policies: how the recommender turns predicted scores into a feed of k items.

Each takes the (users x items) score matrix and returns a (users x k) array of item ids in
feed order (slot 0 at the top). All of them see the same scores; they differ only in how much
they are willing to show items the model is unsure about.
"""
from __future__ import annotations

import numpy as np

from .rng import STREAMS, normals, stream_key, uniforms

POLICIES = ["exploit", "eps", "ucb", "thompson", "mmr", "random"]
LABELS = {"exploit": "Pure exploit", "eps": "ε-greedy", "ucb": "UCB", "thompson": "Thompson", "mmr": "MMR re-rank", "random": "Random"}
DEFAULTS = {"eps": 0.1, "ucb_c": 0.15, "ts_c": 0.3, "mmr_lambda": 0.7}


def topk(adjusted: np.ndarray, k: int) -> np.ndarray:
    """The k highest-scoring items of each row, best first."""
    part = np.argpartition(-adjusted, k - 1, axis=1)[:, :k]
    vals = np.take_along_axis(adjusted, part, axis=1)
    order = np.argsort(-vals, axis=1, kind="stable")
    return np.take_along_axis(part, order, axis=1)


def exploit(pred, ctx):
    return topk(pred, ctx.k)


def ucb(pred, ctx):
    # the bandit index: predicted score plus a bonus that shrinks with the item's exposure
    bonus = ctx.p["ucb_c"] * np.sqrt(np.log(ctx.total_imps + 1.0) / (ctx.item_imps + 1.0))
    return topk(pred + bonus[None, :], ctx.k)


def thompson(pred, ctx):
    # one posterior sample per (user, item): the score plus Gaussian noise whose spread shrinks with exposure
    U, M = pred.shape
    base = (ctx.t - 1) * U * M
    z = normals(ctx.keys["thompson"], np.arange(base, base + U * M)).reshape(U, M)
    return topk(pred + ctx.p["ts_c"] * z / np.sqrt(ctx.item_imps + 1.0)[None, :], ctx.k)


def eps_greedy(pred, ctx, eps=None):
    eps = ctx.p["eps"] if eps is None else eps
    U, M = pred.shape
    k = ctx.k
    base = (ctx.t - 1) * U * k
    coins = uniforms(ctx.keys["eps_coin"], np.arange(base, base + U * k)).reshape(U, k)
    picks = uniforms(ctx.keys["eps_item"], np.arange(base, base + U * k)).reshape(U, k)
    ranked = topk(pred, 2 * k)
    act = ctx.active_items
    na = len(act)
    feeds = np.zeros((U, k), dtype=np.int64)
    for u in range(U):
        chosen = set()
        nxt = 0
        for r in range(k):
            if coins[u, r] < eps:
                j = int(np.floor(picks[u, r] * na))
                while act[j] in chosen:
                    j = (j + 1) % na
                it = int(act[j])
            else:
                while int(ranked[u, nxt]) in chosen:
                    nxt += 1
                it = int(ranked[u, nxt])
                nxt += 1
            chosen.add(it)
            feeds[u, r] = it
    return feeds


def random_feed(pred, ctx):
    return eps_greedy(pred, ctx, eps=1.0)


def mmr(pred, ctx):
    """Maximal marginal relevance: each slot takes the item with the best
    lambda * score - (1 - lambda) * top_score * (its topic is already in the feed),
    where top_score is the user's highest score, so the penalty is on the scores' own scale."""
    lam = ctx.p["mmr_lambda"]
    U, M = pred.shape
    k, T = ctx.k, ctx.n_topics
    present = np.zeros((U, T))
    work = lam * pred.copy()
    penalty = ((1.0 - lam) * pred.max(axis=1))[:, None]
    feeds = np.zeros((U, k), dtype=np.int64)
    rows = np.arange(U)
    for r in range(k):
        adj = work - penalty * present[:, ctx.topic]
        best = adj.argmax(axis=1)
        feeds[:, r] = best
        work[rows, best] = -np.inf
        present[rows, ctx.topic[best]] = 1.0
    return feeds


FUNCS = {"exploit": exploit, "eps": eps_greedy, "ucb": ucb, "thompson": thompson, "mmr": mmr, "random": random_feed}


class Ctx:
    """What a policy may look at besides the scores."""

    def __init__(self, seed, k, topic, params):
        self.k, self.topic, self.p = k, topic, dict(DEFAULTS, **params)
        self.n_topics = int(topic.max()) + 1
        self.keys = {name: stream_key(seed, s) for name, s in STREAMS.items()}
        self.t = 0
        self.item_imps = np.zeros(len(topic))
        self.total_imps = 0.0
        self.active_items = np.arange(len(topic))