"""Exact paired-binary power calculations for the prospective Study 2.

The calculation enumerates the multinomial counts for outcomes where the new
harness alone succeeds, the baseline alone succeeds, or the pair is concordant.
Rejection uses the same two-sided exact McNemar/binomial test preregistered for
the primary analysis.  No experimental outcome data are read by this script.
"""

from __future__ import annotations

import argparse
from math import exp, lgamma


def exact_binomial_pvalue(successes: int, trials: int) -> float:
    """Return scipy-compatible two-sided exact binomial p for p=0.5."""

    if trials == 0:
        return 1.0
    observed = _binomial_probability(successes, trials)
    probability = 0.0
    for value in range(trials + 1):
        candidate = _binomial_probability(value, trials)
        if candidate <= observed + 1e-15:
            probability += candidate
    return min(probability, 1.0)


def _binomial_probability(successes: int, trials: int) -> float:
    log_coefficient = lgamma(trials + 1) - lgamma(successes + 1) - lgamma(
        trials - successes + 1
    )
    return exp(log_coefficient - trials * 0.6931471805599453)


def _multinomial_probability(
    total: int,
    new_only: int,
    baseline_only: int,
    new_only_probability: float,
    baseline_only_probability: float,
) -> float:
    concordant = total - new_only - baseline_only
    concordant_probability = 1.0 - new_only_probability - baseline_only_probability
    counts = (new_only, baseline_only, concordant)
    probabilities = (
        new_only_probability,
        baseline_only_probability,
        concordant_probability,
    )
    log_value = lgamma(total + 1) - sum(lgamma(value + 1) for value in counts)
    for count, probability in zip(counts, probabilities):
        if count and probability == 0.0:
            return 0.0
        if count:
            from math import log

            log_value += count * log(probability)
    return exp(log_value)


def exact_mcnemar_power(
    tasks: int,
    new_only_probability: float,
    baseline_only_probability: float,
    alpha: float = 0.05,
) -> float:
    if tasks <= 0:
        raise ValueError("tasks must be positive")
    if not 0.0 < alpha < 1.0:
        raise ValueError("alpha must lie strictly between zero and one")
    if min(new_only_probability, baseline_only_probability) < 0.0:
        raise ValueError("discordance probabilities cannot be negative")
    if new_only_probability + baseline_only_probability > 1.0:
        raise ValueError("discordance probabilities cannot sum above one")
    power = 0.0
    for new_only in range(tasks + 1):
        for baseline_only in range(tasks - new_only + 1):
            discordant = new_only + baseline_only
            if exact_binomial_pvalue(new_only, discordant) <= alpha:
                power += _multinomial_probability(
                    tasks,
                    new_only,
                    baseline_only,
                    new_only_probability,
                    baseline_only_probability,
                )
    return power


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--tasks", type=int, default=60)
    parser.add_argument("--new-only", type=float, default=0.25)
    parser.add_argument("--baseline-only", type=float, default=0.05)
    parser.add_argument("--alpha", type=float, default=0.05)
    arguments = parser.parse_args()
    power = exact_mcnemar_power(
        arguments.tasks,
        arguments.new_only,
        arguments.baseline_only,
        arguments.alpha,
    )
    print(
        f"tasks={arguments.tasks} new_only={arguments.new_only:.3f} "
        f"baseline_only={arguments.baseline_only:.3f} alpha={arguments.alpha:.4f} "
        f"power={power:.6f}"
    )


if __name__ == "__main__":
    main()