#!/usr/bin/env python3
"""Audit the Dimensional Reality individual-outcome prediction frontier.

The program implements the falsifiable method described in the court memo:

1. Start from a fully specified Born conditional distribution p_t(y | C_t).
2. Let an accessible pre-measurement lift Lambda_t modify it through a
   normalized exponential tilt q_t(y | Lambda_t, C_t).
3. Commit q_t before observing Y_t and update the likelihood-ratio e-process
   E_T = product_t q_t(Y_t) / p_t(Y_t).
4. Combine predeclared candidate lifts with a mixture e-value so model search
   cannot manufacture evidence by selecting the best curve after the fact.
5. Replace a falsely precise Born baseline by a conservative upper envelope
   when calibration leaves a declared composite null.
6. Penalize a countable compiler-generated lift language by prefix-code length.

Under an exact Born null and predictable q_t, E_T is a nonnegative martingale
with mean one. Ville's inequality therefore makes 1/alpha an anytime-valid
rejection threshold. The code verifies the one-step identity, demonstrates
null and synthetic-positive runs, and prints planning sizes. It does not claim
that any real quantum outcome has been predicted.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import random
from statistics import NormalDist


FIVE_SIGMA_ONE_SIDED_ALPHA = 2.87e-7


def normalize(probabilities: list[float]) -> list[float]:
    """Return a strictly positive normalized probability vector."""
    if not probabilities or any(value <= 0 or not math.isfinite(value) for value in probabilities):
        raise ValueError("probabilities must be finite and strictly positive")
    total = sum(probabilities)
    if total <= 0:
        raise ValueError("probability total must be positive")
    return [value / total for value in probabilities]


def exponential_tilt(
    born_probabilities: list[float], lift_scores: list[float]
) -> list[float]:
    """Construct q(y) proportional to p_Born(y) exp(f_y(Lambda, C))."""
    born = normalize(born_probabilities)
    if len(born) != len(lift_scores):
        raise ValueError("Born probabilities and lift scores must have equal length")
    if any(not math.isfinite(score) for score in lift_scores):
        raise ValueError("lift scores must be finite")
    offset = max(lift_scores)
    weights = [p * math.exp(score - offset) for p, score in zip(born, lift_scores)]
    return normalize(weights)


def binary_tilt(born_probability_one: float, lift_score: float) -> float:
    """Binary form: q(1)=sigmoid(logit(p_Born(1)) + lift_score)."""
    if not 0 < born_probability_one < 1:
        raise ValueError("Born probability must lie strictly between zero and one")
    logit = math.log(born_probability_one / (1 - born_probability_one))
    shifted = logit + lift_score
    if shifted >= 0:
        return 1 / (1 + math.exp(-shifted))
    exponential = math.exp(shifted)
    return exponential / (1 + exponential)


def expected_log_gain_under_born(
    born_probabilities: list[float], candidate_probabilities: list[float]
) -> float:
    """Return E_Born[log(q(Y)/p(Y))] = -KL(p || q), never positive."""
    born = normalize(born_probabilities)
    candidate = normalize(candidate_probabilities)
    if len(born) != len(candidate):
        raise ValueError("probability vectors must have equal length")
    return sum(p * math.log(q / p) for p, q in zip(born, candidate))


def expected_e_multiplier_under_born(
    born_probabilities: list[float], candidate_probabilities: list[float]
) -> float:
    """Return E_Born[q(Y)/p(Y)] = 1 for every normalized candidate q."""
    born = normalize(born_probabilities)
    candidate = normalize(candidate_probabilities)
    if len(born) != len(candidate):
        raise ValueError("probability vectors must have equal length")
    return sum(p * q / p for p, q in zip(born, candidate))


def update_log_e(
    log_e_value: float,
    outcome: int,
    born_probabilities: list[float],
    candidate_probabilities: list[float],
) -> float:
    """Update log E after one revealed outcome."""
    born = normalize(born_probabilities)
    candidate = normalize(candidate_probabilities)
    if len(born) != len(candidate):
        raise ValueError("probability vectors must have equal length")
    if not 0 <= outcome < len(born):
        raise ValueError("outcome index is outside the probability vector")
    return log_e_value + math.log(candidate[outcome] / born[outcome])


def validate_upper_envelope(upper_envelope: list[float]) -> list[float]:
    """Validate coordinatewise upper bounds for a composite probability null."""
    if not upper_envelope or any(
        value <= 0 or value > 1 or not math.isfinite(value)
        for value in upper_envelope
    ):
        raise ValueError("null upper bounds must be finite and lie in (0, 1]")
    return list(upper_envelope)


def expected_composite_e_multiplier(
    null_probabilities: list[float],
    upper_envelope: list[float],
    candidate_probabilities: list[float],
) -> float:
    """Return E_p[q(Y)/u(Y)] <= 1 whenever p(y) <= u(y)."""
    null = normalize(null_probabilities)
    envelope = validate_upper_envelope(upper_envelope)
    candidate = normalize(candidate_probabilities)
    if len(null) != len(envelope) or len(null) != len(candidate):
        raise ValueError("null, envelope, and candidate vectors must have equal length")
    if any(p > upper + 1e-15 for p, upper in zip(null, envelope)):
        raise ValueError("the supplied null distribution exceeds its upper envelope")
    return sum(p * q / upper for p, upper, q in zip(null, envelope, candidate))


def update_composite_log_e(
    log_e_value: float,
    outcome: int,
    upper_envelope: list[float],
    candidate_probabilities: list[float],
) -> float:
    """Update log capital against every null distribution below the envelope."""
    envelope = validate_upper_envelope(upper_envelope)
    candidate = normalize(candidate_probabilities)
    if len(envelope) != len(candidate):
        raise ValueError("envelope and candidate vectors must have equal length")
    if not 0 <= outcome < len(envelope):
        raise ValueError("outcome index is outside the probability vector")
    return log_e_value + math.log(candidate[outcome] / envelope[outcome])


def logsumexp(values: list[float]) -> float:
    """Stable logarithm of a sum of exponentials."""
    maximum = max(values)
    return maximum + math.log(sum(math.exp(value - maximum) for value in values))


def mixture_log_e(log_e_values: list[float], weights: list[float]) -> float:
    """Return log(sum_k w_k E_k) for predeclared candidate lifts."""
    if len(log_e_values) != len(weights) or not weights:
        raise ValueError("e-values and mixture weights must have equal length")
    if any(weight < 0 or not math.isfinite(weight) for weight in weights):
        raise ValueError("mixture weights must be finite and nonnegative")
    weight_total = sum(weights)
    if weight_total <= 0:
        raise ValueError("at least one mixture weight must be positive")
    normalized_weights = [weight / weight_total for weight in weights]
    return logsumexp(
        [
            math.log(weight) + log_e
            for weight, log_e in zip(normalized_weights, log_e_values)
            if weight > 0
        ]
    )


def kraft_weights(code_lengths: list[int]) -> tuple[list[float], float]:
    """Return 2^-length model weights and uncommitted cash under Kraft's bound."""
    if not code_lengths or any(
        not isinstance(length, int) or isinstance(length, bool) or length < 0
        for length in code_lengths
    ):
        raise ValueError("code lengths must be nonnegative integers")
    weights = [2.0 ** (-length) for length in code_lengths]
    total = sum(weights)
    if total > 1 + 1e-15:
        raise ValueError("code lengths violate Kraft's inequality")
    return weights, max(0.0, 1 - total)


def kraft_mixture_log_e(log_e_values: list[float], code_lengths: list[int]) -> float:
    """Mix candidate e-values with 2^-length weights and retain unused cash."""
    if len(log_e_values) != len(code_lengths):
        raise ValueError("e-values and code lengths must have equal length")
    weights, cash = kraft_weights(code_lengths)
    terms = [
        math.log(weight) + log_e
        for weight, log_e in zip(weights, log_e_values)
        if weight > 0
    ]
    if cash > 0:
        terms.append(math.log(cash))
    return logsumexp(terms)


def candidate_log_threshold(alpha: float, code_length: int) -> float:
    """Sufficient log E_k for one prefix-coded candidate to cross the mixture."""
    if not 0 < alpha < 1:
        raise ValueError("alpha must lie strictly between zero and one")
    if not isinstance(code_length, int) or isinstance(code_length, bool) or code_length < 0:
        raise ValueError("code length must be a nonnegative integer")
    return math.log(1 / alpha) + code_length * math.log(2)


def mutual_information(joint_probabilities: list[list[float]]) -> float:
    """Return I(Lambda;Y) in nats for a finite joint probability table."""
    if not joint_probabilities or not joint_probabilities[0]:
        raise ValueError("joint table must be nonempty")
    outcome_count = len(joint_probabilities[0])
    if any(len(row) != outcome_count for row in joint_probabilities):
        raise ValueError("joint table must be rectangular")
    flat = [value for row in joint_probabilities for value in row]
    if any(value < 0 or not math.isfinite(value) for value in flat):
        raise ValueError("joint probabilities must be finite and nonnegative")
    total = sum(flat)
    if total <= 0:
        raise ValueError("joint probability total must be positive")
    joint = [[value / total for value in row] for row in joint_probabilities]
    lift_marginal = [sum(row) for row in joint]
    outcome_marginal = [
        sum(row[outcome] for row in joint) for outcome in range(outcome_count)
    ]
    information = 0.0
    for lift, row in enumerate(joint):
        for outcome, probability in enumerate(row):
            if probability > 0:
                information += probability * math.log(
                    probability
                    / (lift_marginal[lift] * outcome_marginal[outcome])
                )
    return information


def coarsen_joint_table(
    joint_probabilities: list[list[float]], group_by_lift: list[int]
) -> list[list[float]]:
    """Push a joint table through a deterministic lift canonicalizer."""
    if len(joint_probabilities) != len(group_by_lift):
        raise ValueError("one group label is required for every lift value")
    if not group_by_lift or any(
        not isinstance(group, int) or isinstance(group, bool) or group < 0
        for group in group_by_lift
    ):
        raise ValueError("group labels must be nonnegative integers")
    outcome_count = len(joint_probabilities[0])
    group_count = max(group_by_lift) + 1
    result = [[0.0] * outcome_count for _ in range(group_count)]
    for group, row in zip(group_by_lift, joint_probabilities):
        if len(row) != outcome_count:
            raise ValueError("joint table must be rectangular")
        for outcome, probability in enumerate(row):
            result[group][outcome] += probability
    return result


def prediction_commitment(
    trial_id: int,
    lift_digest: str,
    probabilities_by_setting: dict[str, list[float]],
    salt: str,
) -> str:
    """Hash a prediction table that can be timestamped before setting/outcome."""
    payload = {
        "trial_id": trial_id,
        "lift_digest": lift_digest,
        "probabilities_by_setting": probabilities_by_setting,
        "salt": salt,
    }
    canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def bernoulli_log_growth_statistics(accuracy: float) -> tuple[float, float]:
    """Mean and variance of log(q/p) for a calibrated balanced predictor."""
    if not 0.5 < accuracy < 1:
        raise ValueError("accuracy must be strictly between 0.5 and 1")
    correct = math.log(2 * accuracy)
    incorrect = math.log(2 * (1 - accuracy))
    mean = accuracy * correct + (1 - accuracy) * incorrect
    variance = (
        accuracy * (correct - mean) ** 2
        + (1 - accuracy) * (incorrect - mean) ** 2
    )
    return mean, variance


def planning_trials(
    accuracy: float,
    alpha: float = FIVE_SIGMA_ONE_SIDED_ALPHA,
    power: float = 0.90,
) -> tuple[int, int]:
    """Return expected-threshold and normal-approximation power trial counts."""
    if not 0 < alpha < 1 or not 0.5 < power < 1:
        raise ValueError("alpha and power must lie in their valid open intervals")
    mean, variance = bernoulli_log_growth_statistics(accuracy)
    threshold = math.log(1 / alpha)
    expected_trials = math.ceil(threshold / mean)
    z_power = NormalDist().inv_cdf(power)
    # Solve N*mean - z_power*sqrt(N*variance) >= threshold for sqrt(N).
    root_n = (
        z_power * math.sqrt(variance)
        + math.sqrt(z_power**2 * variance + 4 * mean * threshold)
    ) / (2 * mean)
    return expected_trials, math.ceil(root_n**2)


def planning_trials_with_code_length(
    accuracy: float,
    code_length: int,
    alpha: float = FIVE_SIGMA_ONE_SIDED_ALPHA,
    power: float = 0.90,
) -> tuple[int, int]:
    """Plan trials after charging one candidate its prefix-code search penalty."""
    mean, variance = bernoulli_log_growth_statistics(accuracy)
    threshold = candidate_log_threshold(alpha, code_length)
    expected_trials = math.ceil(threshold / mean)
    z_power = NormalDist().inv_cdf(power)
    root_n = (
        z_power * math.sqrt(variance)
        + math.sqrt(z_power**2 * variance + 4 * mean * threshold)
    ) / (2 * mean)
    return expected_trials, math.ceil(root_n**2)


def simulate_binary_run(
    trials: int,
    predictive_accuracy: float,
    seed: int,
    born_null: bool,
) -> tuple[float, float]:
    """Simulate a balanced pre-measurement bit lift and return final/max log E."""
    if not 0.5 < predictive_accuracy < 1:
        raise ValueError("predictive accuracy must be strictly between 0.5 and 1")
    generator = random.Random(seed)
    log_e = 0.0
    maximum_log_e = 0.0
    for _ in range(trials):
        lift_bit = generator.randrange(2)
        if born_null:
            outcome = generator.randrange(2)
        else:
            outcome = (
                lift_bit
                if generator.random() < predictive_accuracy
                else 1 - lift_bit
            )
        candidate = [1 - predictive_accuracy, predictive_accuracy]
        if lift_bit == 0:
            candidate.reverse()
        log_e = update_log_e(log_e, outcome, [0.5, 0.5], candidate)
        maximum_log_e = max(maximum_log_e, log_e)
    return log_e, maximum_log_e


def property_audit(trials: int, seed: int) -> tuple[float, float]:
    """Stress composite validity and data processing on random finite tables."""
    if trials < 1:
        raise ValueError("property-audit trials must be positive")
    generator = random.Random(seed)
    maximum_composite_multiplier = 0.0
    maximum_data_processing_excess = -math.inf
    for _ in range(trials):
        outcome_count = generator.randint(2, 6)
        raw_null = [generator.expovariate(1.0) for _ in range(outcome_count)]
        null = normalize(raw_null)
        envelope = [
            p + generator.random() * (1 - p)
            for p in null
        ]
        candidate = normalize(
            [generator.expovariate(1.0) for _ in range(outcome_count)]
        )
        multiplier = expected_composite_e_multiplier(null, envelope, candidate)
        maximum_composite_multiplier = max(maximum_composite_multiplier, multiplier)
        assert multiplier <= 1 + 1e-12

        lift_count = generator.randint(2, 8)
        joint = [
            [generator.expovariate(1.0) for _ in range(outcome_count)]
            for _ in range(lift_count)
        ]
        grouping = [
            generator.randrange(generator.randint(1, lift_count + 1))
            for _ in range(lift_count)
        ]
        full_information = mutual_information(joint)
        coarse_information = mutual_information(coarsen_joint_table(joint, grouping))
        excess = coarse_information - full_information
        maximum_data_processing_excess = max(
            maximum_data_processing_excess, excess
        )
        assert excess <= 1e-12
    return maximum_composite_multiplier, maximum_data_processing_excess


def print_report(trials: int, seed: int) -> None:
    """Run invariant checks and print the frontier audit."""
    born = [0.11, 0.19, 0.23, 0.47]
    candidate = exponential_tilt(born, [-0.4, 0.2, 0.9, -0.1])
    e_multiplier = expected_e_multiplier_under_born(born, candidate)
    log_gain = expected_log_gain_under_born(born, candidate)
    assert math.isclose(e_multiplier, 1.0, abs_tol=1e-14)
    assert log_gain <= 1e-14
    assert exponential_tilt(born, [0.0] * len(born)) == normalize(born)
    assert math.isclose(
        binary_tilt(0.5, math.log(0.51 / 0.49)), 0.51, abs_tol=1e-14
    )

    composite_null = [0.49, 0.51]
    composite_envelope = [0.52, 0.52]
    composite_candidate = [0.48, 0.52]
    composite_multiplier = expected_composite_e_multiplier(
        composite_null, composite_envelope, composite_candidate
    )
    assert composite_multiplier <= 1 + 1e-14
    for null_probability_one in (0.48, 0.50, 0.52):
        multiplier = expected_composite_e_multiplier(
            [1 - null_probability_one, null_probability_one],
            composite_envelope,
            composite_candidate,
        )
        assert multiplier <= 1 + 1e-14

    full_lift_joint = [[0.4, 0.1], [0.1, 0.4]]
    coarsened_joint = coarsen_joint_table(full_lift_joint, [0, 0])
    full_information = mutual_information(full_lift_joint)
    coarsened_information = mutual_information(coarsened_joint)
    assert full_information > 0
    assert math.isclose(coarsened_information, 0.0, abs_tol=1e-14)
    assert coarsened_information <= full_information + 1e-14
    property_trials = min(trials, 5_000)
    maximum_multiplier, maximum_information_excess = property_audit(
        property_trials, seed + 2
    )

    commitment = prediction_commitment(
        trial_id=1,
        lift_digest="example-premeasurement-lift",
        probabilities_by_setting={"X": [0.49, 0.51], "Z": [0.5, 0.5]},
        salt="publish-after-the-run",
    )
    assert len(commitment) == 64

    threshold = math.log(1 / FIVE_SIGMA_ONE_SIDED_ALPHA)
    null_final, null_maximum = simulate_binary_run(
        trials, predictive_accuracy=0.51, seed=seed, born_null=True
    )
    positive_final, positive_maximum = simulate_binary_run(
        trials, predictive_accuracy=0.51, seed=seed + 1, born_null=False
    )

    print("BORN PREDICTION WALL")
    print(f"expected_e_multiplier_under_null,{e_multiplier:.12g}")
    print(f"expected_log_gain_under_null,{log_gain:.12g}")
    print(f"five_sigma_alpha,{FIVE_SIGMA_ONE_SIDED_ALPHA:.12g}")
    print(f"five_sigma_e_threshold,{1 / FIVE_SIGMA_ONE_SIDED_ALPHA:.12g}")
    print(f"five_sigma_log_e_threshold,{threshold:.12g}")
    print(f"example_prediction_commitment,{commitment}")

    print("\nCOMPOSITE BORN NULL")
    print("null_p1,envelope_u0,envelope_u1,candidate_q1,expected_multiplier")
    print(
        f"{composite_null[1]:.6g},{composite_envelope[0]:.6g},"
        f"{composite_envelope[1]:.6g},{composite_candidate[1]:.6g},"
        f"{composite_multiplier:.12g}"
    )

    print("\nPROJECTION DATA-PROCESSING CHECK")
    print("full_lift_mutual_information,coarsened_mutual_information")
    print(f"{full_information:.12g},{coarsened_information:.12g}")
    print(
        f"random_property_trials,{property_trials}\n"
        f"maximum_composite_multiplier,{maximum_multiplier:.12g}\n"
        f"maximum_data_processing_excess,{maximum_information_excess:.12g}"
    )

    print("\nSYNTHETIC SEQUENTIAL RUNS (accuracy claim=0.51)")
    print("world,final_log_e,max_log_e,crossed_five_sigma")
    print(f"Born-null,{null_final:.12g},{null_maximum:.12g},{null_maximum >= threshold}")
    print(
        f"synthetic-accessible-lift,{positive_final:.12g},"
        f"{positive_maximum:.12g},{positive_maximum >= threshold}"
    )

    print("\nFIVE-SIGMA, 90%-POWER PLANNING (balanced binary benchmark)")
    print("accuracy,advantage,mean_log_gain,expected_trials,approx_power_trials")
    for accuracy in (0.55, 0.52, 0.51, 0.505, 0.501):
        mean, _ = bernoulli_log_growth_statistics(accuracy)
        expected_trials, powered_trials = planning_trials(accuracy)
        print(
            f"{accuracy:.6g},{accuracy - 0.5:.6g},{mean:.12g},"
            f"{expected_trials},{powered_trials}"
        )

    print("\nPREFIX-CODE SEARCH PENALTY (accuracy=0.51)")
    print("code_length_bits,required_log_e,expected_trials,approx_power_trials")
    for code_length in (0, 8, 16, 32):
        expected_trials, powered_trials = planning_trials_with_code_length(
            0.51, code_length
        )
        print(
            f"{code_length},{candidate_log_threshold(FIVE_SIGMA_ONE_SIDED_ALPHA, code_length):.12g},"
            f"{expected_trials},{powered_trials}"
        )

    # A convex mixture of valid candidate e-values is itself a valid e-value.
    mixed = mixture_log_e([math.log(0.4), math.log(1.6)], [0.5, 0.5])
    assert math.isclose(math.exp(mixed), 1.0, abs_tol=1e-14)
    assert math.isclose(
        mixture_log_e([math.log(2.0), math.log(7.0)], [1.0, 0.0]),
        math.log(2.0),
        abs_tol=1e-14,
    )
    weights, cash = kraft_weights([2, 3, 3])
    assert math.isclose(sum(weights) + cash, 1.0, abs_tol=1e-14)
    kraft_mixed = kraft_mixture_log_e([0.0, 0.0, 0.0], [2, 3, 3])
    assert math.isclose(kraft_mixed, 0.0, abs_tol=1e-14)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--trials",
        type=int,
        default=120_000,
        help="trials in each illustrative sequential run (default: 120000)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=20_260_915,
        help="deterministic random seed (default: 20260915)",
    )
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    if arguments.trials < 1:
        raise SystemExit("--trials must be positive")
    print_report(arguments.trials, arguments.seed)