"""MSC-P-048 candidate (local, unpublished): reading an A/B test before its planned date.

Standard library only. Synthetic experiment, declared as such: 250 visitors a day, each assigned
at random to one of the two arms with probability one half, as a real A/B test does; a 5 % base
conversion rate; a planned horizon of 28 days and monitoring up to 84. Every count is drawn
exactly from the binomial law by inverting its distribution function, one uniform per draw, from
the declared linear congruential recursion.

Three decision rules are compared: the planned fixed date, stopping at the first significant look,
and a boundary that stays valid however often it is read. That boundary runs on the conversion
clock. Random per-visitor assignment is what makes it legitimate: under the null hypothesis each
conversion belongs to the variant with probability exactly one half, independently of the others,
so the increment of the conversion gap is +1 or -1 at the toss of a coin. Such an increment is
1-sub-Gaussian for every conversion rate, since cosh(l) <= exp(l^2/2) holds for all l. Fixing the
arm sizes instead would make the conditional law hypergeometric and the argument false, which is
why the assignment is drawn here and why an invariant measures it.

The tuning parameter is chosen as the source prescribes: the boundary is tightest at about 8.21
times rho, so rho is searched numerically to minimise the threshold at the target date, and the
ratio is printed as a check against the source's own statement.

The exact size of the fixed-date test is computed, not simulated: for a given control count the
rejection set in the variant count is the outside of a quadratic, so the exact probability follows
from the binomial distribution function. The simulation is then a check, never the source of truth.

Everything is written out so that the R reference prints the same lines.
"""
import math

LCG_SEED = 20260916
LCG_MULTIPLIER = 1103515245
LCG_INCREMENT = 12345
LCG_MODULUS = 2147483648

BASE = 0.05
UPLIFT = 0.015
PER_DAY = 250  # visitors a day in total, each assigned at random to one arm
PLANNED_DAYS = 28
MONITOR_DAYS = 84
ALPHA = 0.05
Z_FIXED = 1.959963985
LOOK_POINTS = (7, 14, 28, 84)
RHO_DAYS = (7, 28, 56)  # calibration dates declared a priori, beside the source optimum
REPLICATIONS = 20000
EXAMPLE_LIMIT = 200
GRID_STEPS = 2000  # grid on which the sub-Gaussian inequality is checked
PROBE_DAYS = 2000  # days drawn by the invariants to measure the assignment and the coin flip
WALKS = 5000
WALK_STEPS = 2000
GUARD_REPLICATIONS = 2000  # replications of the blocking null-hypothesis check
STICKY_REPLICATIONS = 5000  # replications of the sticky-assignment measurement
STICKY_BLOCK = 10  # visitors sharing one arm when the assignment is drawn but sticky


class Declared:
    """Declared recursion; uniforms are (state + 0.5) / modulus, never 0 or 1."""

    def __init__(self, seed):
        self.state = seed

    def uniform(self):
        self.state = (LCG_MULTIPLIER * self.state + LCG_INCREMENT) % LCG_MODULUS
        return (self.state + 0.5) / LCG_MODULUS


def first_terms(p, limit):
    """(1 - p) to the n for n = 0 to limit, by repeated multiplication.

    The same operation in both languages, rather than a power, whose last bit differs between
    Python and R. Built once per rate, so a draw costs one uniform and a short recursion.
    """
    terms = [1.0]
    for _ in range(limit):
        terms.append(terms[-1] * (1.0 - p))
    return terms


TERMS_HALF = first_terms(0.5, PER_DAY)
TERMS_BASE = first_terms(BASE, PER_DAY)
TERMS_UPLIFT = first_terms(BASE + UPLIFT, PER_DAY)


def binomial(stream, n, p, terms):
    """Exact binomial draw by inversion: one uniform, then the probability recursion."""
    u = stream.uniform()
    ratio = p / (1.0 - p)
    term = terms[n]
    total = term
    k = 0
    while u > total and k < n:
        term *= (n - k) / (k + 1.0) * ratio
        k += 1
        total += term
    return k


def pooled_z(n_control, control, n_variant, variant):
    """Two-sided pooled z statistic of two proportions, arms of possibly different sizes."""
    if n_control <= 0 or n_variant <= 0:
        return 0.0
    pooled = (control + variant) / (n_control + n_variant)
    if pooled <= 0.0 or pooled >= 1.0:
        return 0.0
    spread = math.sqrt(pooled * (1.0 - pooled) * (1.0 / n_control + 1.0 / n_variant))
    return (variant / n_variant - control / n_control) / spread


def mixture_bound(conversions, rho):
    """Two-sided normal-mixture boundary on the conversion gap, read on the conversion clock.

    Closed form of the two-sided normal mixture, equation (14) of Howard, Ramdas, McAuliffe and
    Sekhon (2018), with the initial likelihood ratio l0 = 1: crossed at any time with probability
    at most ALPHA. Proposition 6 of the same article is the one-sided form and is not used here.
    """
    shifted = conversions + rho
    return math.sqrt(shifted * math.log(shifted / (rho * ALPHA * ALPHA)))


def tightest_rho(target):
    """Rho that minimises the threshold at a target number of conversions, by ternary search.

    The source states the boundary is tightest at a fixed multiple of rho, its equation (21); the
    search reproduces that multiple instead of hard-coding it, and an invariant checks the ratio.
    """
    low, high = 0.1, float(target)
    for _ in range(200):
        first = low + (high - low) / 3.0
        second = high - (high - low) / 3.0
        if mixture_bound(target, first) < mixture_bound(target, second):
            high = second
        else:
            low = first
    return (low + high) / 2.0


def conversions_expected(day):
    """Conversions expected by a date, under the declared base rate and daily traffic."""
    return PER_DAY * day * BASE


def draw_sticky_day(stream, uplift):
    """A day where the assignment is drawn but sticky: blocks of visitors share one arm.

    This is what cookie-level or session-level assignment does in a real tool. The shares stay at
    one half on average, so the measured invariants pass, but the conversion labels are no longer
    independent and the boundary loses its guarantee. The programs measure how much.
    """
    n_variant = 0
    blocks = PER_DAY // STICKY_BLOCK
    for _ in range(blocks):
        if stream.uniform() < 0.5:
            n_variant += STICKY_BLOCK
    n_control = PER_DAY - n_variant
    variant = binomial(stream, n_variant, BASE + uplift, TERMS_UPLIFT if uplift else TERMS_BASE)
    control = binomial(stream, n_control, BASE, TERMS_BASE)
    return n_control, control, n_variant, variant


def draw_day(stream, uplift):
    """One day: the split of visitors between arms, then the conversions of each arm."""
    n_variant = binomial(stream, PER_DAY, 0.5, TERMS_HALF)
    n_control = PER_DAY - n_variant
    variant = binomial(stream, n_variant, BASE + uplift, TERMS_UPLIFT if uplift else TERMS_BASE)
    control = binomial(stream, n_control, BASE, TERMS_BASE)
    return n_control, control, n_variant, variant


def check_invariants(stream):
    """Checks run before anything is computed; any failure stops the run.

    Two of them are measured on the drawn data, the assignment and the conversion coin flip,
    and can fail. The sub-Gaussian inequality is a mathematical identity, verified numerically
    rather than assumed; its least favourable point is the first of the grid.
    """
    failures = []
    worst = None
    for i in range(1, GRID_STEPS + 1):
        lam = i * 0.01
        gap = math.log(math.cosh(lam)) - lam * lam / 2.0
        worst = gap if worst is None or gap > worst else worst
    if worst > 0.0:
        failures.append("sub_gaussian_inequality")

    visitors = 0
    to_variant = 0
    conversions = 0
    variant_conversions = 0
    for _ in range(PROBE_DAYS):
        n_control, control, n_variant, variant = draw_day(stream, 0.0)
        visitors += n_control + n_variant
        to_variant += n_variant
        conversions += control + variant
        variant_conversions += variant
    assignment = to_variant / visitors
    coin = variant_conversions / conversions
    if abs(assignment - 0.5) > 0.005:
        failures.append("balanced_assignment")
    if abs(coin - 0.5) > 0.01:
        failures.append("conversion_coin_flip")

    target = conversions_expected(PLANNED_DAYS)
    rho = tightest_rho(target)
    ratio = target / rho if rho > 0.0 else 0.0
    if rho <= 0.0:
        failures.append("positive_rho")
    if abs(ratio - 8.212) > 0.01:
        failures.append("tightest_point_matches_source")
    for probe in (1.0, 10.0, 100.0, 500.0, 2000.0, target, 10.0 * target):
        shifted = probe + rho
        rearranged = shifted * (math.log(shifted) - math.log(rho) - 2.0 * math.log(ALPHA))
        if abs(mixture_bound(probe, rho) ** 2 - rearranged) > 1e-6:
            failures.append("boundary_closed_form")
            break

    probe_stream = Declared(LCG_SEED)
    before = probe_stream.state
    binomial(probe_stream, PER_DAY, BASE, TERMS_BASE)
    if probe_stream.state != (LCG_MULTIPLIER * before + LCG_INCREMENT) % LCG_MODULUS:
        failures.append("one_uniform_per_draw")

    # The clock the boundary is actually read on, checked end to end on a short test: the verdicts
    # must be reproducible from the conversion total alone, recomputed outside run_test.
    clock_stream = Declared(LCG_SEED)
    looks = run_test(clock_stream, 30, 0.0, [rho])
    control = 0
    variant = 0
    for day, (_, gap, crossed) in enumerate(looks, start=1):
        pass
    replay = Declared(LCG_SEED)
    control = 0
    variant = 0
    for day, (_, gap, crossed) in enumerate(looks, start=1):
        day_n_control, day_control, day_n_variant, day_variant = draw_day(replay, 0.0)
        control += day_control
        variant += day_variant
        if gap != variant - control:
            failures.append("boundary_reads_conversion_gap")
            break
        expected = abs(variant - control) >= mixture_bound(control + variant, rho) if control + variant > 0 else False
        if crossed[0] != expected:
            failures.append("boundary_reads_conversion_clock")
            break

    print("metric.invariant.sub_gaussian_worst_gap=%.3e grid=%d_values_lambda_0.01_to_%.2f"
          % (worst, GRID_STEPS, GRID_STEPS * 0.01))
    print("metric.invariant.measured visitors=%d to_variant_share=%.5f conversions=%d variant_share=%.5f"
          % (visitors, assignment, conversions, coin))
    print("metric.invariant.tightest_point target_over_rho=%.3f source_states=8.212" % ratio)
    print("metric.invariant.checks_failed=%d %s" % (len(failures), ",".join(failures) if failures else "none"))
    if failures:
        raise SystemExit("invariant failed: " + ",".join(failures))


def boundary_crossing_check(stream, walks, steps, rho):
    """Independent check of the boundary: a Rademacher walk, where the null holds by construction."""
    crossed = 0
    for _ in range(walks):
        position = 0
        for step in range(1, steps + 1):
            position += 1 if stream.uniform() < 0.5 else -1
            if abs(position) >= mixture_bound(step, rho):
                crossed += 1
                break
    return crossed


def log_choose(n, k):
    return math.lgamma(n + 1.0) - math.lgamma(k + 1.0) - math.lgamma(n - k + 1.0)


def exact_size(n):
    """Exact two-sided size of the pooled z test with n per arm, with no simulation.

    For a fixed control count c the test rejects when 2n (v - c)^2 >= z^2 (c + v)(2n - c - v), a
    quadratic in v whose outside is the rejection set; the exact probability then follows from the
    distribution function. Read at equal arm sizes, which is the design in expectation.
    """
    pmf = [math.exp(log_choose(n, k) + k * math.log(BASE) + (n - k) * math.log(1.0 - BASE))
           for k in range(n + 1)]
    cdf = []
    running = 0.0
    for value in pmf:
        running += value
        cdf.append(running)
    k = Z_FIXED * Z_FIXED
    size = 0.0
    a = 2.0 * n + k
    for c in range(n + 1):
        b = -2.0 * (2.0 * n * c + k * n - k * c)
        constant = 2.0 * n * c * c - k * (2.0 * n * c - c * c)
        disc = b * b - 4.0 * a * constant
        if disc <= 0.0:
            continue
        root = math.sqrt(disc)
        below = math.floor((-b - root) / (2.0 * a))
        above = math.ceil((-b + root) / (2.0 * a))
        mass = 0.0
        if below >= 0.0:
            mass += cdf[int(below) if below <= n else n]
            if c == 0:
                mass -= pmf[0]  # the pair (0, 0) is never rejected: the statistic is zero
        if above <= n:
            index = int(above) if above > 0.0 else 0
            mass += 1.0 - (cdf[index - 1] if index > 0 else 0.0)
        size += pmf[c] * mass
    return size


def run_test(stream, days, uplift, rhos):
    """One experiment read every day: z statistic, conversion gap, and each boundary's verdict."""
    control = 0
    variant = 0
    n_control = 0
    n_variant = 0
    looks = []
    for _ in range(days):
        day_n_control, day_control, day_n_variant, day_variant = draw_day(stream, uplift)
        n_control += day_n_control
        control += day_control
        n_variant += day_n_variant
        variant += day_variant
        gap = variant - control
        conversions = control + variant
        crossed = tuple(abs(gap) >= mixture_bound(conversions, rho) if conversions > 0 else False
                        for rho in rhos)
        looks.append((pooled_z(n_control, control, n_variant, variant), gap, crossed))
    return looks


def summary(hits, total):
    share = hits / total
    return hits, share, math.sqrt(max(share * (1.0 - share), 0.0) / total)


def first_day(flags, limit):
    for day, flag in enumerate(flags, start=1):
        if day > limit:
            return 0
        if flag:
            return day
    return 0


def main():
    stream = Declared(LCG_SEED)
    check_invariants(stream)

    target = conversions_expected(PLANNED_DAYS)
    labels = ["optimal"] + ["day_%d" % day for day in RHO_DAYS]
    rhos = [tightest_rho(target)] + [conversions_expected(day) for day in RHO_DAYS]
    reference = 0  # the source-optimal tuning is this page's reference

    print("metric.design.base_rate=%.3f uplift_synthetic=%.3f visitors_per_day=%d assignment=random_half"
          % (BASE, UPLIFT, PER_DAY))
    print("metric.design.planned_days=%d monitored_days=%d alpha=%.3f replications=%d"
          % (PLANNED_DAYS, MONITOR_DAYS, ALPHA, REPLICATIONS))
    print("metric.design.visitors_expected_per_arm_at_planned_day=%.1f conversions_expected=%.1f"
          % (PER_DAY * PLANNED_DAYS / 2.0, target))
    print("metric.design.z_fixed=%.6f boundary_clock=conversions" % Z_FIXED)
    for label, rho in zip(labels, rhos):
        over_root = mixture_bound(target, rho) / math.sqrt(target)
        print("metric.design.boundary tuning=%s rho=%.4f threshold_over_sqrt_conversions=%.4f threshold_in_z=%.4f"
              % (label, rho, over_root, over_root / math.sqrt(1.0 - BASE)))

    crossed = boundary_crossing_check(Declared(LCG_SEED), WALKS, WALK_STEPS, rhos[reference])
    walk_rate = crossed / WALKS
    print("metric.check.boundary_on_rademacher_walk walks=%d steps=%d crossed=%d rate=%.5f alpha=%.3f verdict=%s"
          % (WALKS, WALK_STEPS, crossed, walk_rate, ALPHA, "pass" if walk_rate <= ALPHA else "fail"))
    if walk_rate > ALPHA:
        raise SystemExit("boundary check failed: crossing rate above alpha")

    guard = 0
    guard_stream = Declared(LCG_SEED)
    for _ in range(GUARD_REPLICATIONS):
        looks = run_test(guard_stream, MONITOR_DAYS, 0.0, [rhos[reference]])
        if first_day([look[2][0] for look in looks], MONITOR_DAYS):
            guard += 1
    guard_rate = guard / GUARD_REPLICATIONS
    print("metric.check.null_false_alarm replications=%d looks=%d crossed=%d rate=%.5f alpha=%.3f verdict=%s"
          % (GUARD_REPLICATIONS, MONITOR_DAYS, guard, guard_rate, ALPHA, "pass" if guard_rate <= ALPHA else "fail"))
    if guard_rate > ALPHA:
        raise SystemExit("null false-alarm check failed: rate above alpha")

    sticky = 0
    sticky_stream = Declared(LCG_SEED)
    for _ in range(STICKY_REPLICATIONS):
        control = 0
        variant = 0
        crossed_once = False
        for _day in range(MONITOR_DAYS):
            _, day_control, _, day_variant = draw_sticky_day(sticky_stream, 0.0)
            control += day_control
            variant += day_variant
            conversions = control + variant
            if conversions > 0 and abs(variant - control) >= mixture_bound(conversions, rhos[reference]):
                crossed_once = True
                break
        if crossed_once:
            sticky += 1
    print("metric.check.sticky_assignment block_visitors=%d replications=%d crossed=%d rate=%.5f alpha=%.3f"
          % (STICKY_BLOCK, STICKY_REPLICATIONS, sticky, sticky / STICKY_REPLICATIONS, ALPHA))

    for day in LOOK_POINTS:
        n = PER_DAY * day // 2
        print("metric.exact.fixed_date_size horizon_days=%d n_per_arm=%d size=%.6f"
              % (day, n, exact_size(n)))

    for attempt in range(1, EXAMPLE_LIMIT + 1):
        looks = run_test(stream, PLANNED_DAYS, 0.0, rhos)
        day = first_day([abs(z) >= Z_FIXED for z, _, _ in looks], PLANNED_DAYS)
        if day:
            stop = looks[day - 1]
            end = looks[PLANNED_DAYS - 1]
            print("metric.example.no_effect_declared_by_peeking attempt=%d day=%d z=%.4f gap=%d"
                  % (attempt, day, stop[0], stop[1]))
            print("metric.example.same_test_at_planned_day day=%d z=%.4f gap=%d significant=%s"
                  % (PLANNED_DAYS, end[0], end[1], "yes" if abs(end[0]) >= Z_FIXED else "no"))
            print("metric.example.same_test_boundary_crossed_by_planned_day=%s"
                  % ("yes" if any(look[2][reference] for look in looks) else "no"))
            break

    for label, uplift in (("no_effect", 0.0), ("true_uplift", UPLIFT)):
        fixed = {limit: 0 for limit in LOOK_POINTS}
        peeks = {limit: 0 for limit in LOOK_POINTS}
        weekly = 0
        after_first_week = 0
        crossed_any = [0] * len(rhos)
        crossed_by_day = [[0] * (MONITOR_DAYS + 1) for _ in rhos]
        peek_days = []
        boundary_days = []
        for _ in range(REPLICATIONS):
            looks = run_test(stream, MONITOR_DAYS, uplift, rhos)
            significant = [abs(z) >= Z_FIXED for z, _, _ in looks]
            for limit in LOOK_POINTS:
                if significant[limit - 1]:
                    fixed[limit] += 1
                day = first_day(significant, limit)
                if day:
                    peeks[limit] += 1
                    if limit == PLANNED_DAYS:
                        peek_days.append(day)
            if any(significant[day - 1] for day in (7, 14, 21, PLANNED_DAYS)):
                weekly += 1
            if any(significant[day - 1] for day in range(7, PLANNED_DAYS + 1)):
                after_first_week += 1
            for index in range(len(rhos)):
                day = first_day([look[2][index] for look in looks], MONITOR_DAYS)
                if day:
                    crossed_any[index] += 1
                    for later in range(day, MONITOR_DAYS + 1):
                        crossed_by_day[index][later] += 1
                    if index == reference:
                        boundary_days.append(day)

        for limit in LOOK_POINTS:
            hits, share, mcse = summary(fixed[limit], REPLICATIONS)
            print("metric.%s.fixed_date horizon_days=%d declared=%d of %d rate=%.5f mcse=%.5f"
                  % (label, limit, hits, REPLICATIONS, share, mcse))
            hits, share, mcse = summary(peeks[limit], REPLICATIONS)
            print("metric.%s.daily_peeking horizon_days=%d looks=%d declared=%d of %d rate=%.5f mcse=%.5f"
                  % (label, limit, limit, hits, REPLICATIONS, share, mcse))
        hits, share, mcse = summary(weekly, REPLICATIONS)
        print("metric.%s.weekly_reading horizon_days=%d looks=4 declared=%d of %d rate=%.5f mcse=%.5f"
              % (label, PLANNED_DAYS, hits, REPLICATIONS, share, mcse))
        hits, share, mcse = summary(after_first_week, REPLICATIONS)
        print("metric.%s.daily_from_day_7 horizon_days=%d looks=22 declared=%d of %d rate=%.5f mcse=%.5f"
              % (label, PLANNED_DAYS, hits, REPLICATIONS, share, mcse))
        for index, tuning in enumerate(labels):
            hits, share, mcse = summary(crossed_any[index], REPLICATIONS)
            print("metric.%s.boundary tuning=%s looks=%d declared=%d of %d rate=%.5f mcse=%.5f"
                  % (label, tuning, MONITOR_DAYS, hits, REPLICATIONS, share, mcse))
            hits, share, mcse = summary(crossed_by_day[index][PLANNED_DAYS], REPLICATIONS)
            print("metric.%s.boundary tuning=%s by_planned_day declared=%d of %d rate=%.5f mcse=%.5f"
                  % (label, tuning, hits, REPLICATIONS, share, mcse))
        if label == "true_uplift":
            goal = fixed[PLANNED_DAYS] / REPLICATIONS
            for index, tuning in enumerate(labels):
                reached = next((d for d in range(1, MONITOR_DAYS + 1)
                                if crossed_by_day[index][d] / REPLICATIONS >= goal), 0)
                before = (crossed_by_day[index][reached - 1] / REPLICATIONS) if reached > 1 else 0.0
                print("metric.true_uplift.boundary tuning=%s day_matching_fixed_power=%d ratio=%.2f goal=%.5f power_day_before=%.5f"
                      % (tuning, reached, (reached / PLANNED_DAYS) if reached else 0.0, goal, before))
            if peek_days:
                peek_days.sort()
                print("metric.true_uplift.peeking_median_stop_day=%d within=%d_days"
                      % (peek_days[len(peek_days) // 2], PLANNED_DAYS))
            if boundary_days:
                boundary_days.sort()
                print("metric.true_uplift.boundary_median_stop_day=%d within=%d_days"
                      % (boundary_days[len(boundary_days) // 2], MONITOR_DAYS))


if __name__ == "__main__":
    main()