"""Chapter 17 lab, part 1: test-time compute (Tasks 1-3).

Spend more samples at inference time, then pick one answer. How much does it
help, and which way of picking is best? numpy and the standard library only.

Tests: test_ttc.py.   Run: python -m pytest -q test_ttc.py
"""

from collections import Counter, defaultdict  # noqa: F401  (handy for Task 3)
from math import comb  # noqa: F401  (handy for Tasks 1-2)

import numpy as np  # noqa: F401


# ---------------------------------------------------------------- Task 1
def pass_at_n(p, n):
    """pass@N: the chance that at least one of n independent samples is correct,
    when each is correct with probability p. (The chapter's "oracle ceiling".)

        pass@N = 1 - (1 - p)^n

    Example: p = 0.3, n = 4  ->  1 - 0.7^4 = 0.7599.
    Works for floats or numpy arrays of p.
    """
    # TODO: one line.
    raise NotImplementedError("Task 1: implement pass_at_n")


def pass_at_k(n, c, k):
    """Unbiased estimate of pass@k from n samples of which c were correct
    (Chen et al. 2021, the Codex paper). With k <= n:

        pass@k_hat = 1 - C(n - c, k) / C(n, k)

    = 1 - P(a random size-k subset of the n samples contains no correct one).
    If n - c < k every subset contains a correct sample, so return 1.0.
    Use math.comb (exact integers); return a float.

    Example: n = 4 samples, c = 1 correct, k = 2  ->  1 - C(3,2)/C(4,2) = 0.5.
    """
    # TODO: handle n - c < k, then the formula.
    raise NotImplementedError("Task 1: implement pass_at_k")


# ---------------------------------------------------------------- Task 2
def majority_vote_accuracy(p, n):
    """Exact accuracy of majority vote with BINARY answers (right / wrong):
    n independent samples, each right with probability p. A tie (possible
    when n is even) is broken by a fair coin.

        sum over c (number of right samples) of  C(n, c) p^c (1 - p)^(n - c) * w(c)
        where w(c) = 1 if 2c > n,  0.5 if 2c == n,  0 otherwise.

    Example: p = 0.6, n = 3  ->  0.6^3 + 3 * 0.6^2 * 0.4 = 0.648.
    """
    # TODO: loop over c = 0..n.
    raise NotImplementedError("Task 2: implement majority_vote_accuracy")


# ---------------------------------------------------------------- Task 3
def majority_vote(answers):
    """Self-consistency: return the most common answer in the list.

    Ties: return the tied answer that appears FIRST in `answers`.
    Example: majority_vote([7, 3, 3, 7, 5]) -> 7   (7 and 3 tie; 7 came first)
    """
    # TODO: count (collections.Counter), then scan `answers` in order.
    raise NotImplementedError("Task 3: implement majority_vote")


def best_of_n(answers, scores):
    """Best-of-N: return the answer of the single highest-scoring sample.

    answers : sequence of length N.   scores : sequence of N floats (verifier scores).
    Ties: the earliest sample wins (np.argmax already does this).
    Example: best_of_n([7, 3, 3], [0.2, 0.9, 0.1]) -> 3
    """
    # TODO: one line.
    raise NotImplementedError("Task 3: implement best_of_n")


def weighted_vote(answers, scores):
    """Weighted majority vote ("best-of-N weighted"): add up the verifier
    scores of all samples that give the same answer; return the answer with
    the largest total.

    Ties: the tied answer that appears first in `answers`.
    Example: weighted_vote([7, 3, 3], [0.8, 0.5, 0.4]) -> 3   (3: 0.9 > 7: 0.8)
    """
    # TODO: accumulate totals per answer (a dict), then scan `answers` in order.
    raise NotImplementedError("Task 3: implement weighted_vote")