%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

# Import Bayesian utility functions
from bayesian import (
    test_non_inferiority_weakly_informative,
    select_best_variant
)

# Import plotting utilities
from plotting_utils import (
    plot_weakly_informative_prior_with_variants,
    plot_multiple_posteriors_comparison
)

# Control group data
n_control = 32106  # Number of users
x_control = 22772  # Number who completed
control_rate = x_control / n_control

print(f"Control Group:")
print(f"  Sample size: {n_control:,}")
print(f"  Conversions: {x_control:,}")
print(f"  Conversion rate: {control_rate:.2%}")

# Variant data
variants_data = {
    'A': {'n': 4625, 'x': 3244},
    'B': {'n': 2100, 'x': 1433},
    'C': {'n': 2022, 'x': 1396}
}

print(f"\nVariants:")
for name, data in variants_data.items():
    rate = data['x'] / data['n']
    print(f"  {name}: n={data['n']:3d}, x={data['x']:3d}, rate={rate:.2%}")

# Test parameters
epsilon = 0.02  # 2% non-inferiority margin (acceptable degradation, business decision)
print(f"\nNon-inferiority margin (ε): {epsilon:.1%}")
print(f"Non-inferiority threshold: {control_rate - epsilon:.2%}")

Control Group:
  Sample size: 32,106
  Conversions: 22,772
  Conversion rate: 70.93%

Variants:
  A: n=4625, x=3244, rate=70.14%
  B: n=2100, x=1433, rate=68.24%
  C: n=2022, x=1396, rate=69.04%

Non-inferiority margin (ε): 2.0%
Non-inferiority threshold: 68.93%

# Domain knowledge: Guestimate that adding passkey creation (2 extra clicks) will degrade by ~2%, 
# we use it to center our prior so this is not a threshold at all, just a prior belief which will change after seeing data. 
# we can play with the value but it has to ve be less than epsilon otherwise it woudl mean we would
# know even before testing that this is inferiro and there woud be no point testing.
expected_degradation = 0.01

# Run non-inferiority test
results = test_non_inferiority_weakly_informative(
    n_control=n_control,
    x_control=x_control,
    variants_data=variants_data,
    epsilon=epsilon,  # Business: can tolerate 5% degradation
    expected_degradation=expected_degradation,  # Domain: expect 2% degradation
    alpha_prior_strength=20,  # Weak prior (high entropy)
    threshold=0.95  # 95% probability required to be sure, can adjust based on risk tolerance
)

print("="*80)
print("PRIOR AND THRESHOLD SETUP")
print("="*80)
print(f"Control rate: {control_rate:.2%}")
print(f"Expected degradation (domain knowledge): {expected_degradation:.1%}")
print(f"  → Prior centered at: {control_rate - expected_degradation:.2%}")
print(f"Maximum acceptable degradation (business): {epsilon:.1%}")
print(f"  → Test threshold at: {control_rate - epsilon:.2%}")
print(f"\nThis means:")
print(f"  • Prior says: 'I expect variant around {control_rate - expected_degradation:.1%}'")
print(f"  • Test says: 'Must be above {control_rate - epsilon:.1%} to pass'")

# Display results
print("\n" + "="*80)
print("NON-INFERIORITY TEST RESULTS")
print("="*80)

for variant_name, result in results.items():
    status = "✓ NON-INFERIOR" if result['is_non_inferior'] else "✗ NOT NON-INFERIOR"
    print(f"\nVariant {variant_name}: {status}")
    print(f"  P(variant > threshold): {result['probability']:.2%}")
    print(f"  Posterior mean: {result['variant_rate']:.2%}")
    print(f"  Prior mean: {result['prior_mean']:.2%}")
    print(f"  Observed rate: {variants_data[variant_name]['x']/variants_data[variant_name]['n']:.2%}")

# Summary
non_inferior_count = sum(1 for r in results.values() if r['is_non_inferior'])
print(f"\n{'='*80}")
print(f"Summary: {non_inferior_count}/{len(variants_data)} variants are non-inferior")
print(f"{'='*80}")

================================================================================
PRIOR AND THRESHOLD SETUP
================================================================================
Control rate: 70.93%
Expected degradation (domain knowledge): 1.0%
  → Prior centered at: 69.93%
Maximum acceptable degradation (business): 2.0%
  → Test threshold at: 68.93%

This means:
  • Prior says: 'I expect variant around 69.9%'
  • Test says: 'Must be above 68.9% to pass'

================================================================================
NON-INFERIORITY TEST RESULTS
================================================================================

Variant A: ✓ NON-INFERIOR
  P(variant > threshold): 96.38%
  Posterior mean: 70.14%
  Prior mean: 69.93%
  Observed rate: 70.14%

Variant B: ✗ NOT NON-INFERIOR
  P(variant > threshold): 25.54%
  Posterior mean: 68.26%
  Prior mean: 69.93%
  Observed rate: 68.24%

Variant C: ✗ NOT NON-INFERIOR
  P(variant > threshold): 55.12%
  Posterior mean: 69.05%
  Prior mean: 69.93%
  Observed rate: 69.04%

================================================================================
Summary: 1/3 variants are non-inferior
================================================================================

# Create visualization - simplified usage!
fig, ax = plot_weakly_informative_prior_with_variants(results)
plt.show()

# Select best variant among all (or filter to non-inferior only)
# For this example, we'll analyze all variants
selection_results = select_best_variant(
    variants_data=variants_data,
    alpha_prior=1,  # Non-informative prior for selection
    beta_prior=1,
    credible_level=0.95,
    n_simulations=100000
)

# Display results
print("="*80)
print("BEST VARIANT SELECTION")
print("="*80)

print(f"\nProbability each variant is best:")
for name, prob in selection_results['probabilities'].items():
    bar = '█' * int(prob * 60)
    print(f"  {name}: {prob:.2%} {bar}")

winner = selection_results['best_variant']
winner_prob = selection_results['probabilities'][winner]
print(f"\n{'='*80}")
print(f"WINNER: Variant {winner}")
print(f"  Probability of being best: {winner_prob:.2%}")
print(f"  Posterior mean: {selection_results['posterior_means'][winner]:.2%}")
print(f"  95% Credible interval: [{selection_results['credible_intervals'][winner][0]:.2%}, {selection_results['credible_intervals'][winner][1]:.2%}]")
print(f"  Expected loss: {selection_results['expected_loss'][winner]:.4f}")
print(f"{'='*80}")

================================================================================
BEST VARIANT SELECTION
================================================================================

Probability each variant is best:
  A: 78.08% ██████████████████████████████████████████████
  B: 4.29% ██
  C: 17.63% ██████████

================================================================================
WINNER: Variant A
  Probability of being best: 78.08%
  Posterior mean: 70.13%
  95% Credible interval: [68.81%, 71.44%]
  Expected loss: 0.0014
================================================================================

# Prepare posteriors for plotting
from scipy.stats import beta as beta_dist

posteriors = {}
for name, data in variants_data.items():
    # Using non-informative prior Beta(1,1) for fair comparison
    alpha_post = data['x'] + 1
    beta_post = data['n'] - data['x'] + 1
    
    posteriors[name] = {
        'alpha': alpha_post,
        'beta': beta_post,
        'mean': alpha_post / (alpha_post + beta_post),
        'ci_95': (
            beta_dist.ppf(0.025, alpha_post, beta_post),
            beta_dist.ppf(0.975, alpha_post, beta_post)
        )
    }

# Create comparison plot
fig, ax = plot_multiple_posteriors_comparison(
    posteriors=posteriors,
    control_group_conversion_rate=control_rate,
    epsilon=epsilon
)
plt.show()

"""
Bayesian utilities for A/B testing.

This module contains Bayesian methods for non-inferiority testing,
variant selection, and conversion rate analysis using Beta-Bernoulli
conjugate models with Monte Carlo simulation.
"""

import numpy as np
from scipy.stats import beta as beta_dist


def test_non_inferiority(
    n_control,
    x_control,
    variants_data,
    epsilon,
    alpha_prior,
    beta_prior,
    threshold=0.95,
):
    """
    Test non-inferiority of multiple variants against a control.

    Uses Bayesian Beta-Bernoulli conjugate model with Monte Carlo simulation
    to compute the probability that each variant's conversion rate is within
    an acceptable degradation margin of the control.

    Parameters
    ----------
    n_control : int
        Number of samples in control group
    x_control : int
        Number of successes in control group
    variants_data : dict
        Dictionary with variant names as keys and {'n': samples, 'x': successes} as values
        Example: {'A': {'n': 1000, 'x': 200}, 'B': {'n': 1000, 'x': 215}}
    epsilon : float
        Non-inferiority margin (e.g., 0.03 for 3%)
    alpha_prior : float
        Alpha parameter for Beta prior
    beta_prior : float
        Beta parameter for Beta prior
    threshold : float, optional
        Probability threshold for declaring non-inferiority (default: 0.95)

    Returns
    -------
    dict : Dictionary with results for each variant containing:
        - 'is_non_inferior': bool, whether variant is non-inferior
        - 'probability': float, P(variant > control - epsilon)
        - 'control_rate': float, posterior mean of control
        - 'variant_rate': float, posterior mean of variant
        - 'posterior_params': tuple, (alpha, beta) of variant posterior

    """
    # Control posterior
    alpha_control = x_control + alpha_prior
    beta_control = n_control - x_control + beta_prior
    control_rate = alpha_control / (alpha_control + beta_control)

    # Boundary for non-inferiority
    boundary = control_rate - epsilon

    results = {}
    n_simulations = 100000

    # Sample from control posterior once (reuse for all variants)
    control_samples = beta_dist.rvs(alpha_control, beta_control, size=n_simulations)

    for variant_name, data in variants_data.items():
        # Variant posterior
        alpha_variant = data["x"] + alpha_prior
        beta_variant = data["n"] - data["x"] + beta_prior
        variant_rate = alpha_variant / (alpha_variant + beta_variant)

        # Sample from variant posterior
        variant_samples = beta_dist.rvs(alpha_variant, beta_variant, size=n_simulations)

        # Compute P(variant > control - epsilon)
        prob_non_inferior = np.mean(variant_samples > (control_samples - epsilon))

        results[variant_name] = {
            "is_non_inferior": prob_non_inferior >= threshold,
            "probability": prob_non_inferior,
            "control_rate": control_rate,
            "variant_rate": variant_rate,
            "posterior_params": (alpha_variant, beta_variant),
        }

    return results


def test_non_inferiority_weakly_informative(
    n_control,
    x_control,
    variants_data,
    epsilon,
    expected_degradation=None,
    alpha_prior_strength=20,
    threshold=0.95,
):
    """
    Test non-inferiority using weakly informative prior based on domain knowledge.

    This function constructs a weakly informative prior centered at your expected
    variant performance (based on domain knowledge), then tests against a separate
    non-inferiority threshold (based on business requirements).

    Key insight: Prior belief (where you expect the variant to be) is SEPARATE from
    the decision threshold (worst acceptable performance).

    The prior is constructed as:
    - α_prior = alpha_prior_strength (default: 20, for high entropy/wide uncertainty)
    - β_prior = (α_prior / target_mean) - α_prior
    - target_mean = control_rate - expected_degradation

    The test threshold is:
    - non_inferiority_threshold = control_rate - epsilon

    Parameters
    ----------
    n_control : int
        Number of samples in control group
    x_control : int
        Number of successes in control group
    variants_data : dict
        Dictionary with variant names as keys and {'n': samples, 'x': successes} as values
        Example: {'A': {'n': 561, 'x': 381}, 'B': {'n': 285, 'x': 192}}
    epsilon : float
        Non-inferiority margin - maximum acceptable degradation (business requirement)
        Example: 0.05 means "can tolerate up to 5% degradation"
    expected_degradation : float, optional
        Expected degradation based on domain knowledge (e.g., "adding 2 clicks will
        degrade by ~2%"). If None, defaults to epsilon (conservative).
        Should typically be LESS than epsilon.
        Example: 0.02 means "expect 2% degradation"
    alpha_prior_strength : float, optional
        Strength parameter for the prior (default: 20). Smaller values give
        wider (more uncertain) priors. Typical values: 10-30.
    threshold : float, optional
        Probability threshold for declaring non-inferiority (default: 0.95)
        Variant is non-inferior if P(variant > control - epsilon) >= threshold

    Returns
    -------
    dict : Dictionary with results for each variant containing:
        - 'is_non_inferior': bool, whether variant is non-inferior
        - 'probability': float, P(variant > control - epsilon)
        - 'control_rate': float, observed control conversion rate
        - 'variant_rate': float, posterior mean of variant
        - 'posterior_params': tuple, (alpha, beta) of variant posterior
        - 'prior_params': tuple, (alpha_prior, beta_prior) used
        - 'prior_mean': float, mean of the prior distribution

    """
    # Compute control conversion rate
    control_rate = x_control / n_control

    # Determine expected degradation (defaults to epsilon if not specified)
    if expected_degradation is None:
        expected_degradation = epsilon

    # Construct weakly informative prior centered at expected performance
    # This reflects domain knowledge, separate from the business threshold
    target_prior_mean = control_rate - expected_degradation
    alpha_prior = alpha_prior_strength
    beta_prior = (alpha_prior / target_prior_mean) - alpha_prior

    # Verify prior is valid (must have positive parameters)
    if beta_prior <= 0:
        raise ValueError(
            f"Invalid prior parameters: beta_prior={beta_prior:.4f} <= 0. "
            f"This can happen when epsilon is too large relative to control_rate. "
            f"Try reducing epsilon or increasing alpha_prior_strength."
        )

    results = {}

    for variant_name, data in variants_data.items():
        # Variant posterior: Beta(x + α_prior, n - x + β_prior)
        n = data["n"]
        x = data["x"]
        alpha_posterior = x + alpha_prior
        beta_posterior = (n - x) + beta_prior
        variant_posterior_mean = alpha_posterior / (alpha_posterior + beta_posterior)

        # Non-inferiority threshold
        non_inferiority_threshold = control_rate - epsilon

        # Direct calculation: P(variant > threshold) using Beta CDF
        # P(X > threshold) = 1 - P(X <= threshold) = 1 - CDF(threshold)
        prob_non_inferior = 1 - beta_dist.cdf(
            non_inferiority_threshold, alpha_posterior, beta_posterior
        )

        results[variant_name] = {
            "is_non_inferior": prob_non_inferior >= threshold,
            "probability": prob_non_inferior,
            "control_rate": control_rate,
            "variant_rate": variant_posterior_mean,
            "posterior_params": (alpha_posterior, beta_posterior),
            "prior_params": (alpha_prior, beta_prior),
            "prior_mean": target_prior_mean,
            "threshold": non_inferiority_threshold,  # Store the actual test threshold
            "epsilon": epsilon,  # Store epsilon for reference
            "n": n,  # Store sample size for plotting
            "x": x,  # Store successes for plotting
        }

    return results


def select_best_variant(
    variants_data,
    alpha_prior=1,
    beta_prior=1,
    credible_level=0.95,
    n_simulations=100000,
):
    """
    Select the best variant among multiple options using Bayesian approach.

    Uses Monte Carlo simulation to determine which variant has the highest
    probability of being the best, along with expected loss calculations
    for decision-making under uncertainty.

    Parameters
    ----------
    variants_data : dict
        Dictionary with variant names as keys and {'n': samples, 'x': successes} as values
        Example: {'A': {'n': 800, 'x': 168}, 'B': {'n': 800, 'x': 172}}
    alpha_prior : float, optional
        Alpha parameter for Beta prior (default: 1 for uniform)
    beta_prior : float, optional
        Beta parameter for Beta prior (default: 1 for uniform)
    credible_level : float, optional
        Credible interval level (default: 0.95)
    n_simulations : int, optional
        Number of Monte Carlo simulations (default: 100000)

    Returns
    -------
    dict : Dictionary containing:
        - 'best_variant': str, name of variant most likely to be best
        - 'probabilities': dict, P(each variant is best)
        - 'posterior_means': dict, posterior mean for each variant
        - 'credible_intervals': dict, (lower, upper) credible interval for each variant
        - 'expected_loss': dict, expected loss from choosing each variant

    Examples
    --------
    >>> variants = {
    ...     'A': {'n': 800, 'x': 168},
    ...     'B': {'n': 800, 'x': 172},
    ...     'C': {'n': 800, 'x': 165}
    ... }
    >>> result = select_best_variant(variants)
    >>> print(result['best_variant'])
    'B'
    >>> print(result['probabilities'])
    {'A': 0.31, 'B': 0.47, 'C': 0.22}
    """
    variant_names = list(variants_data.keys())
    posteriors = {}
    samples = {}

    # Compute posteriors and draw samples
    for name, data in variants_data.items():
        alpha_post = data["x"] + alpha_prior
        beta_post = data["n"] - data["x"] + beta_prior

        posteriors[name] = {
            "alpha": alpha_post,
            "beta": beta_post,
            "mean": alpha_post / (alpha_post + beta_post),
        }

        # Draw samples from posterior
        samples[name] = beta_dist.rvs(alpha_post, beta_post, size=n_simulations)

        # Compute credible interval
        ci_lower = beta_dist.ppf((1 - credible_level) / 2, alpha_post, beta_post)
        ci_upper = beta_dist.ppf(1 - (1 - credible_level) / 2, alpha_post, beta_post)
        posteriors[name]["credible_interval"] = (ci_lower, ci_upper)

    # Monte Carlo: count how often each variant is best
    best_counts = {name: 0 for name in variant_names}

    for i in range(n_simulations):
        # Get samples for this iteration
        sample_values = {name: samples[name][i] for name in variant_names}

        # Find best variant in this simulation
        best_variant = max(sample_values, key=sample_values.get)
        best_counts[best_variant] += 1

    # Calculate probabilities
    probabilities = {name: count / n_simulations for name, count in best_counts.items()}

    # Expected loss: E[max(all) - this variant]
    expected_loss = {}
    for name in variant_names:
        max_samples = np.maximum.reduce([samples[v] for v in variant_names])
        losses = max_samples - samples[name]
        expected_loss[name] = np.mean(losses)

    # Determine best variant
    best_variant = max(probabilities, key=probabilities.get)

    return {
        "best_variant": best_variant,
        "probabilities": probabilities,
        "posterior_means": {name: posteriors[name]["mean"] for name in variant_names},
        "credible_intervals": {
            name: posteriors[name]["credible_interval"] for name in variant_names
        },
        "expected_loss": expected_loss,
    }

Bayesian A/B Test Workflow¶

Workflow Overview¶

Key Advantages of Bayesian Approach¶

1. Setup: Import Libraries and Define Data¶

2. Define Experiment Data¶

3. Non-Inferiority Test¶

Key Insight: Domain Knowledge vs Business Tolerance¶

4. Visualize Non-Inferiority Test¶

5. Select Best Variant¶

6. Visualize Variant Comparison¶

Summary¶

Key Takeaway¶

Utility Functions Used in the Notebook¶

Appendix: Mathematical Foundations¶

A.1 The Beta Distribution¶

Definition¶

Probability Density Function (PDF)¶

Properties¶

Intuition: Pseudo-Counts¶

A.2 The Bernoulli Likelihood¶

Single Trial¶

Multiple Independent Trials¶

A.3 Bayesian Inference: From Prior to Posterior¶

Bayes' Theorem¶

The Evidence (Marginal Likelihood)¶

Proportionality Form¶

A.4 Beta-Binomial Conjugacy: The Closed-Form Miracle¶

Why Conjugacy Matters¶

The Derivation¶

The Result: Bayesian Update Rule¶

A.5 Why the Integration Works Out¶

The Normalizing Constant¶

Explicit Integration¶

The Full Posterior¶

Why This is Remarkable¶

A.6 Example: Updating with Real Data¶

Scenario¶

Bayesian Update¶

Interpretation¶

A.7 Sequential Updates: The Bayesian Learning Loop¶

Conjugacy Enables Sequential Learning¶

Mathematical Equivalence¶

A.8 Summary: The Mathematical Beauty of Beta-Binomial Conjugacy¶

The Full Picture¶

Why This Matters for A/B Testing¶

The Fundamental Insight¶