"""The loop: show feeds, simulate clicks, retrain, repeat; with the metrics of each round."""
from __future__ import annotations

import numpy as np

from . import policies
from .recommender import ALS
from .rng import STREAMS, stream_key, uniforms
from .world import TOPICS, World, oracle

METRICS = ["diversity", "coverage", "gini", "utility", "clicks", "bubble", "top1", "new_item"]


def position_bias(k: int) -> np.ndarray:
    """How much less a slot gets looked at, by position: 1 / log2(rank + 2), as in DCG."""
    return 1.0 / np.log2(np.arange(k) + 2.0)


def gini(x: np.ndarray) -> float:
    x = np.sort(np.asarray(x, dtype=float))
    n = len(x)
    s = x.sum()
    if s <= 0:
        return 0.0
    return float((2.0 * np.sum(np.arange(1, n + 1) * x)) / (n * s) - (n + 1.0) / n)


def run(world: World, policy: str = "exploit", rounds: int = 30, params: dict | None = None,
        switch_round: int = 0, switch_policy: str | None = None, new_item_round: int = 0,
        visitor_ratings: dict | None = None, repeat_decay: float = 0.7, system_discount: bool = True,
        factors: int = 6, reg: float = 1.0, sweeps: int = 2, watch: int | None = None, on_round=None) -> dict:
    """Run `rounds` rounds of the feedback loop and return the metrics per round.

    Metrics are over the simulated users only (the visitor, if any, is the last user and is
    left out because their taste is a guess). `new_item_round` releases the held-back item
    (world.new_item) at that round. `switch_round` changes to `switch_policy` from that round on.
    The feed and clicks of user `watch` (the visitor by default) are recorded every round.
    """
    params = params or {}
    U, M, k, T = world.n_users, world.n_items, world.k, len(TOPICS)
    n_sim = U - (1 if visitor_ratings is not None else 0)
    seed = world.seed
    model = ALS(U, M, factors, reg, sweeps, seed)
    ctx = policies.Ctx(seed, k, world.topic, params)
    ctx.n_topics = T
    active = np.ones(M, dtype=bool)
    if world.new_item >= 0 and new_item_round > 0:
        active[world.new_item] = False
    ctx.active_items = np.nonzero(active)[0]
    if visitor_ratings:
        v = U - 1
        items = np.array([int(i) for i in visitor_ratings], dtype=np.int64)
        liked = np.array([bool(visitor_ratings[i]) for i in visitor_ratings])
        model.observe(np.full(len(items), v), items, liked)

    key_click = stream_key(seed, STREAMS["click"])
    jitter = 1e-6 * uniforms(stream_key(seed, STREAMS["jitter"]), np.arange(U * M)).reshape(U, M)
    pos = position_bias(k)
    best = oracle(world, active)
    sim = slice(0, n_sim)
    rows = np.arange(U)[:, None]
    out = {m: [] for m in METRICS}
    if watch is None and visitor_ratings is not None:
        watch = U - 1
    watched = {"user": watch, "feeds": [], "clicks": []}
    new_first = None
    n_top1 = max(1, int(np.ceil(0.01 * M)))

    for t in range(1, rounds + 1):
        if world.new_item >= 0 and t == new_item_round:
            active[world.new_item] = True
            ctx.active_items = np.nonzero(active)[0]
            best = oracle(world, active)
        pol = switch_policy if (switch_round and t >= switch_round and switch_policy) else policy
        ctx.t = t

        # 1. rank: what the system believes, plus the policy's appetite for the unknown
        pred = model.scores()
        if system_discount:
            pred = pred * np.power(repeat_decay, model.N)
        pred = pred + jitter
        pred[:, ~active] = -1e9
        feeds = policies.FUNCS[pol](pred, ctx)

        # 2. clicks: true affinity x position bias x boredom with repeats, against a common coin
        shown_before = model.N[rows, feeds]
        p = world.affinity[rows, feeds] * pos[None, :] * np.power(repeat_decay, shown_before)
        base = (t - 1) * U * k
        coins = uniforms(key_click, np.arange(base, base + U * k)).reshape(U, k)
        clicked = coins < p

        # 3. metrics, over the simulated users
        f = feeds[sim]
        imps = np.bincount(f.ravel(), minlength=M).astype(float)
        imps_active = imps[active]
        topics = world.topic[f]
        distinct = np.array([len(set(row)) for row in topics])
        out["diversity"].append(float(distinct.mean()))
        out["coverage"].append(float((imps_active > 0).mean()))
        out["gini"].append(gini(imps_active))
        out["utility"].append(float((world.affinity[rows[sim], f].mean(axis=1) / best[sim]).mean()))
        out["clicks"].append(float(clicked[sim].sum(axis=1).mean()))
        out["bubble"].append(float((topics == world.top_topic[sim][:, None]).mean()))
        out["top1"].append(float(np.sort(imps_active)[-n_top1:].sum() / imps_active.sum()))
        if world.new_item >= 0:
            ni = float(imps[world.new_item])
            out["new_item"].append(ni)
            if ni > 0 and new_first is None:
                new_first = t
        else:
            out["new_item"].append(0.0)
        if watch is not None:
            watched["feeds"].append([int(i) for i in feeds[watch]])
            watched["clicks"].append([bool(c) for c in clicked[watch]])

        # 4. learn from what was shown, and only from that
        model.observe(np.repeat(np.arange(U), k), feeds.ravel(), clicked.ravel())
        ctx.item_imps += np.bincount(feeds.ravel(), minlength=M)
        ctx.total_imps += U * k
        model.fit()
        if on_round:
            on_round(t, {m: out[m][-1] for m in METRICS})

    return {"policy": policy, "rounds": rounds, "switch_round": switch_round, "switch_policy": switch_policy,
            "new_item_round": new_item_round, "params": ctx.p, "metrics": out, "new_item_first_round": new_first,
            "watched": watched if watch is not None else None,
            "item_imps": [float(x) for x in ctx.item_imps]}