"""Secret Santa from a hat with redraws: how often does the draw get stuck?

People draw in order. If you pull a name you may not have (your own, or with
the couples rule your partner's), you put it back and draw again, so in effect
you get a uniform pick among the allowed names still in the hat. The draw is
stuck when someone finds only forbidden names left.
"""

from fractions import Fraction
from functools import lru_cache

from ..rng import Stream

NAME = "hatdraw"
FAMILY = "Secret Santa"
EXACT_MAX = 14

RULES = """`people` people (numbered 0..people-1) put their names in a hat and then draw one at a time,
in order 0, 1, 2, ... Each person draws a name; if it is a name they are not allowed to have,
they put it back and draw again, so in effect they take a uniformly random name among the
allowed names still in the hat. Names taken are not returned. A person may never take their
own name. If `couples` is true, the people form couples (0,1), (2,3), (4,5), ... and a person
may also never take their partner's name. The draw FAILS if some person, at their turn, finds
that every name left in the hat is one they are not allowed to have. Report the probability
that the draw fails."""
AGENT_TESTS = [{"people": 8, "couples": False}, {"people": 10, "couples": True}]
AGENT_RETURN = "p"


def pages():
    out = [{"params": {"people": n, "couples": False}, "chunks": 1_000, "chunk_size": 10_000} for n in range(3, 31)]
    out += [{"params": {"people": n, "couples": True}, "chunks": 1_000, "chunk_size": 10_000} for n in range(4, 21, 2)]
    return out


def slug(p):
    return f"secret-santa-hat-draw-gets-stuck-{p['people']}-people" + ("-couples" if p["couples"] else "")


def title(p):
    if p["couples"]:
        return (f"Secret Santa with {p['people'] // 2} couples: what is the chance the hat draw gets stuck "
                "when partners can't draw each other?")
    return f"Secret Santa with {p['people']} people: what is the chance the hat draw gets stuck?"


def intro(p):
    base = (f"{p['people']} people draw names from a hat in turn. Anyone who pulls their own name"
            + (" or their partner's" if p["couples"] else "")
            + " puts it back and draws again. How often does someone reach the hat and find only names they can't take?")
    return base


def columns(p):
    return ["games", "stuck", "stuck_last", "stuck_earlier"]


def sim_chunk(p, seed, chunk, n):
    m, couples = p["people"], p["couples"]
    st = Stream(seed, chunk)
    below = st.below
    out = [n, 0, 0, 0]
    for _ in range(n):
        hat = list(range(m))
        for i in range(m):
            partner = i ^ 1 if couples else -1
            cnt = 0
            for s in hat:
                if s != i and s != partner:
                    cnt += 1
            if cnt == 0:
                out[1] += 1
                if i == m - 1:
                    out[2] += 1
                else:
                    out[3] += 1
                break
            r = below(cnt)
            for idx, s in enumerate(hat):
                if s != i and s != partner:
                    if r == 0:
                        hat.pop(idx)
                        break
                    r -= 1
    return out


def exact(p):
    m, couples = p["people"], p["couples"]
    if m > EXACT_MAX:
        return None

    @lru_cache(maxsize=None)
    def f(i, mask):
        """(P stuck at the last person, P stuck earlier) from person i with `mask` in the hat."""
        if i == m:
            return Fraction(0), Fraction(0)
        partner = i ^ 1 if couples else -1
        allowed = [s for s in range(m) if mask >> s & 1 and s != i and s != partner]
        if not allowed:
            return (Fraction(1), Fraction(0)) if i == m - 1 else (Fraction(0), Fraction(1))
        last = early = Fraction(0)
        w = Fraction(1, len(allowed))
        for s in allowed:
            a, b = f(i + 1, mask & ~(1 << s))
            last += w * a
            early += w * b
        return last, early

    last, early = f(0, (1 << m) - 1)
    return {"headline": last + early, "cols": {"stuck": last + early, "stuck_last": last, "stuck_earlier": early},
            "method": "every possible sequence of draws, weighted by its probability (exact recursion)"}


def headline(p, tot):
    return tot[1], tot[0]


def row_labels(p):
    return [("stuck", "The draw gets stuck"), ("stuck_last", "… at the last person"),
            ("stuck_earlier", "… before the last person")]