"""n of a kind with n dice in r rolls, keeping the most common face."""

from fractions import Fraction
from itertools import product

from ..rng import Stream

NAME = "dice"
FAMILY = "Dice"
WORD = {1: "one", 2: "two", 3: "three", 4: "four", 5: "five", 6: "six"}

RULES = """Roll `dice` fair six-sided dice. If they all show the same face, you have won.
Otherwise keep every die showing your most common face (if two faces tie for most common,
keep either one; it does not change the odds) and re-roll all the others. You may roll at
most `rolls` times in total (the first roll counts as one). Report the probability that
at some point all `dice` dice show the same face."""
AGENT_TESTS = [{"dice": 5, "rolls": 3}, {"dice": 4, "rolls": 2}]
AGENT_RETURN = "p"


def pages():
    out = []
    for d in (3, 4, 5, 6):
        for r in (1, 2, 3):
            big = (d, r) == (5, 3)
            out.append({"params": {"dice": d, "rolls": r},
                        "chunks": 10_000 if big else 1_000, "chunk_size": 100_000 if big else 20_000})
    return out


def slug(p):
    d, r = p["dice"], p["rolls"]
    return f"{WORD[d]}-of-a-kind-with-{d}-dice-in-{WORD[r]}-roll{'s' if r > 1 else ''}"


def title(p):
    d, r = p["dice"], p["rolls"]
    return f"What is the chance of {WORD[d]} of a kind with {d} dice in {WORD[r]} roll{'s' if r > 1 else ''}?"


def intro(p):
    d, r = p["dice"], p["rolls"]
    if r == 1:
        return f"Roll {d} dice once. How often do all {d} show the same face?"
    return (f"Roll {d} dice. Keep every die showing your most common face and re-roll the rest, "
            f"up to {WORD[r]} rolls in all. How often do you end with all {d} the same?")


def columns(p):
    d, r = p["dice"], p["rolls"]
    return ["games"] + [f"best_{k}" for k in range(1, d)] + [f"all_roll_{j}" for j in range(1, r + 1)]


def sim_chunk(p, seed, chunk, n):
    d, rolls = p["dice"], p["rolls"]
    st = Stream(seed, chunk)
    out = [0] * (1 + (d - 1) + rolls)
    out[0] = n
    below = st.below
    for _ in range(n):
        dice = []
        roll = 1
        need = d
        while True:
            for _ in range(need):
                dice.append(below(6))
            counts = [0] * 6
            for x in dice:
                counts[x] += 1
            best = 0
            for f in range(1, 6):
                if counts[f] > counts[best]:
                    best = f
            k = counts[best]
            if k == d:
                out[d - 1 + roll] += 1
                break
            if roll == rolls:
                out[k] += 1
                break
            dice = [best] * k
            need = d - k
            roll += 1
    return out


def exact(p):
    """Markov chain on the size of the largest group, by enumerating every roll."""
    d, rolls = p["dice"], p["rolls"]

    def dist(keep):
        m = d - keep
        out = [Fraction(0)] * (d + 1)
        for o in product(range(6), repeat=m):
            c = [o.count(f) for f in range(6)]
            c[0] += keep
            out[max(c)] += Fraction(1, 6 ** m)
        return out

    T = {k: dist(k) for k in range(1, d)}
    cur = dist(0)
    cols = {"all_roll_1": cur[d]}
    for j in range(2, rolls + 1):
        nxt = [Fraction(0)] * (d + 1)
        for k in range(1, d):
            for m in range(1, d + 1):
                nxt[m] += cur[k] * T[k][m]
        cols[f"all_roll_{j}"] = nxt[d]
        cur = nxt
    for k in range(1, d):
        cols[f"best_{k}"] = cur[k]
    head = sum(cols[f"all_roll_{j}"] for j in range(1, rolls + 1))
    return {"headline": head, "cols": cols, "method": "Markov chain over every possible roll"}


def headline(p, tot):
    d, r = p["dice"], p["rolls"]
    return sum(tot[d:d + r]) , tot[0]   # successes, games


def row_labels(p):
    d, r = p["dice"], p["rolls"]
    rows = [(f"all_roll_{j}", f"All {d} the same on roll {j}") for j in range(1, r + 1)]
    rows += [(f"best_{k}", f"Ended with a best group of {k}") for k in range(d - 1, 0, -1)]
    return rows