"""How many packs to complete a set of n items, with s distinct items per pack."""

from functools import lru_cache
from math import comb

from ..rng import Stream

NAME = "collector"
FAMILY = "Collecting"

RULES = """A set has `items` different items, all equally common. Each pack contains `per_pack` items,
all different from each other within the pack, chosen uniformly at random among all such
combinations (packs are independent of each other). You buy packs one at a time until you
own every item at least once. Report the expected number of packs bought ("mean") and the
sample standard deviation of the number of packs ("sd")."""
AGENT_TESTS = [{"items": 50, "per_pack": 1}, {"items": 100, "per_pack": 5}]
AGENT_RETURN = "mean"

SETS = [(10, 1, 10_000), (20, 1, 10_000), (50, 1, 5_000), (100, 1, 2_000), (250, 1, 1_000),
        (300, 5, 1_000), (500, 5, 1_000), (680, 5, 1_000), (700, 7, 1_000)]


def pages():
    out = []
    for n, s, size in SETS:
        p = {"items": n, "per_pack": s}
        q = _markov(n, s)["quantiles"]
        p.update({"k50": q[0], "k90": q[1], "k99": q[2]})
        out.append({"params": p, "chunks": 1_000, "chunk_size": size})
    return out


def slug(p):
    if p["per_pack"] == 1:
        return f"how-many-packs-to-complete-a-set-of-{p['items']}"
    return f"how-many-packs-to-complete-a-{p['items']}-sticker-album-{p['per_pack']}-per-pack"


def title(p):
    if p["per_pack"] == 1:
        return f"How many tries to collect all {p['items']} items, one at a time?"
    return f"How many packs to fill a {p['items']}-sticker album with {p['per_pack']} stickers per pack?"


def intro(p):
    if p["per_pack"] == 1:
        return (f"Each try gives you one of {p['items']} equally common items at random. "
                "How many tries until you have them all?")
    return (f"Each pack holds {p['per_pack']} different stickers out of {p['items']}, all equally common, "
            "and you trade with nobody. How many packs until the album is full?")


def columns(p):
    return ["games", "sum_packs", "sum_packs_sq", "min_packs", "max_packs",
            f"done_by_{p['k50']}", f"done_by_{p['k90']}", f"done_by_{p['k99']}"]


def sim_chunk(p, seed, chunk, n):
    m, s = p["items"], p["per_pack"]
    ks = (p["k50"], p["k90"], p["k99"])
    st = Stream(seed, chunk)
    below = st.below
    tot = tot2 = 0
    lo, hi = None, 0
    done = [0, 0, 0]
    for _ in range(n):
        have = bytearray(m)
        got = 0
        t = 0
        while got < m:
            t += 1
            pack = []
            while len(pack) < s:
                x = below(m)
                if x in pack:
                    continue
                pack.append(x)
            for x in pack:
                if not have[x]:
                    have[x] = 1
                    got += 1
        tot += t
        tot2 += t * t
        lo = t if lo is None or t < lo else lo
        hi = t if t > hi else hi
        for j in range(3):
            if t <= ks[j]:
                done[j] += 1
    return [n, tot, tot2, lo, hi] + done


@lru_cache(maxsize=None)
def _markov(n, s):
    """Distribution of the number of packs, by a Markov chain on how many items you own.

    From m owned, a pack brings j new items with probability C(n-m, j) C(m, s-j) / C(n, s).
    Double precision; the tail is followed until it is below 1e-15.
    """
    c = comb(n, s)
    P = [[comb(n - m, j) * comb(m, s - j) / c for j in range(s + 1)] for m in range(n + 1)]
    dist = [0.0] * (n + 1)
    dist[0] = 1.0
    k = 0
    mean = 0.0
    second = 0.0
    cdf = {}
    quant = [None, None, None]
    targets = (0.5, 0.9, 0.99)
    while True:
        tail = sum(dist[:n])  # summed directly: 1 - dist[n] can stall on rounding
        mean += tail
        second += (2 * k + 1) * tail
        if (tail < 1e-15 and all(quant)) or k > 200_000:
            break
        k += 1
        new = [0.0] * (n + 1)
        new[n] = dist[n]
        for m in range(n):
            pm = dist[m]
            if pm == 0.0:
                continue
            row = P[m]
            for j in range(s + 1):
                if row[j]:
                    new[m + j] += pm * row[j]
        dist = new
        cdf[k] = dist[n]
        for i, t in enumerate(targets):
            if quant[i] is None and dist[n] >= t:
                quant[i] = k
    return {"mean": mean, "sd": (second - mean * mean) ** 0.5, "quantiles": quant, "cdf": cdf}


def exact(p):
    r = _markov(p["items"], p["per_pack"])
    cols = {f"done_by_{p[k]}": r["cdf"][p[k]] for k in ("k50", "k90", "k99") if k in p}
    return {"headline": r["mean"], "sd": r["sd"], "cols": cols, "float": True,
            "method": "Markov chain on how many items you own, in double precision"}


def headline(p, tot):
    return tot[1], tot[0]   # sum of packs, games -> mean


def row_labels(p):
    return [(f"done_by_{p['k50']}", f"Done within {p['k50']} packs"),
            (f"done_by_{p['k90']}", f"Done within {p['k90']} packs"),
            (f"done_by_{p['k99']}", f"Done within {p['k99']} packs")]