# The task and the decision rule, both fixed before data collection.
# Task: each round the asset's factor is 1.3 or 0.8 with equal probability; leverage L
# (0 to 3, in steps of 0.25) turns it into 1 + 0.3 L or 1 - 0.2 L.
import math
g = lambda L: 0.5 * math.log(1 + 0.3 * L) + 0.5 * math.log(1 - 0.2 * L)
grid = [i / 4 for i in range(13)]
print("leverage  expected gain/round  time-average growth/round")
for L in grid[::2]:
    print(f"{L:8.2f} {5 * L:+19.1f} % {100 * (math.exp(g(L)) - 1):+24.2f} %")
best = max(grid, key=g)
print(f"best leverage on the grid: {best} (exact Kelly 5/6 = {5 / 6:.3f});"
      f" expected gain rises up to L = 3")
# Share of players below their start after 40 rounds (exact, from the binomial count of
# up-rounds): what the ergodic display shows at each leverage.
for L in (5 / 6, 1.0, 2.0, 3.0):
    below = sum(math.comb(40, k) for k in range(41)
                if k * math.log(1 + 0.3 * L) + (40 - k) * math.log(1 - 0.2 * L) < 0) / 2 ** 40
    print(f"L = {L:.2f}: {below:.1%} of players end the 40 rounds below their start")

# Decision rule. d = arm B (ergodic display) minus arm A (arithmetic display) in mean
# chosen leverage, standardised, with its 95 % CI. Lower leverage is the predicted way.
SESOI = 0.30
def verdict(d, lo, hi, comprehension_ok=True):
    if not comprehension_ok:
        return f"{d:+.2f}: NOT INTERPRETABLE - over 15 % failed comprehension in an arm"
    if hi <= -SESOI:  v = "B lowers leverage by at least the smallest effect of interest"
    elif hi < 0:      v = "B lowers leverage; size not shown to reach 0.30 SD"
    elif lo > 0:      v = "B RAISES leverage (opposite to prediction)"
    elif lo >= -SESOI and hi <= SESOI: v = "no difference worth acting on"
    else:             v = "unclear"
    return f"{d:+.2f} ({lo:+.2f} to {hi:+.2f}): {v}"
cases = ((-0.55, -0.73, -0.37), (-0.45, -0.63, -0.27), (-0.15, -0.33, 0.03),
         (-0.20, -0.36, -0.04), (0.05, -0.13, 0.23), (0.25, 0.07, 0.43))
for case in cases:
    print(verdict(*case))
print("Secondary, reported with intervals, never decisive: share of rounds above 2 x Kelly;")
print("distance |L - L*|, which also catches over-caution; final wealth.")
# leverage  expected gain/round  time-average growth/round
#     0.00                +0.0 %                    +0.00 %
#     0.50                +2.5 %                    +1.73 %
#     1.00                +5.0 %                    +1.98 %
#     1.50                +7.5 %                    +0.75 %
#     2.00               +10.0 %                    -2.02 %
#     2.50               +12.5 %                    -6.46 %
#     3.00               +15.0 %                   -12.82 %
# best leverage on the grid: 0.75 (exact Kelly 5/6 = 0.833); expected gain rises up to L = 3
# L = 0.83: 21.5% of players end the 40 rounds below their start
# L = 1.00: 31.8% of players end the 40 rounds below their start
# L = 2.00: 56.3% of players end the 40 rounds below their start
# L = 3.00: 86.6% of players end the 40 rounds below their start
# -0.55 (-0.73 to -0.37): B lowers leverage by at least the smallest effect of interest
# -0.45 (-0.63 to -0.27): B lowers leverage; size not shown to reach 0.30 SD
# -0.15 (-0.33 to +0.03): unclear
# -0.20 (-0.36 to -0.04): B lowers leverage; size not shown to reach 0.30 SD
# +0.05 (-0.13 to +0.23): no difference worth acting on
# +0.25 (+0.07 to +0.43): B RAISES leverage (opposite to prediction)
# Secondary, reported with intervals, never decisive: share of rounds above 2 x Kelly;
# distance |L - L*|, which also catches over-caution; final wealth.

# Sample size for the pre-registered randomised behavioural experiment. Participants play
# the same repeated multiplicative investment task (40 rounds; each round they choose a
# leverage from 0 to 3 on a risky asset; bonus proportional to final wealth). Arm A sees
# the asset described by its arithmetic expected return and the AVERAGE final wealth of
# many simulated players; arm B sees the same parameters as the MEDIAN path, the
# time-average growth rate and the share of simulated players ending below the start.
# Primary outcome: mean chosen leverage over the 40 rounds. Two-sided alpha 0.05.
import math
from statistics import NormalDist
N = NormalDist()

def n_per_arm(d, alpha=0.05, power=0.90):
    # normal-approximation formula, rounded up, plus 1 to allow for using the t distribution
    return math.ceil(2 * (N.inv_cdf(1 - alpha / 2) + N.inv_cdf(power)) ** 2 / d ** 2) + 1

for d in (0.20, 0.25, 0.30, 0.40):
    print(f"standardised effect {d:.2f}: {n_per_arm(d, power=0.8):4d} per arm for 80 % power,"
          f" {n_per_arm(d):4d} for 90 %")
n = n_per_arm(0.30)
recruit = math.ceil(n / 0.85)
print(f"chosen: 0.30 SD, 90 % power -> {n} analysable per arm")
print(f"recruit {recruit} per arm allowing 15 % exclusions (comprehension checks),"
      f" {2 * recruit} in all")
# Baseline leverage (mean of 5 practice rounds, before any display appears) enters the
# analysis as a covariate. If it correlates r = 0.5 with the outcome, power rises; the
# sample is still sized without it, leaving a margin.
r = 0.5
pw = N.cdf(0.30 / math.sqrt(2 * (1 - r ** 2) / n) - N.inv_cdf(0.975))
print(f"with a baseline covariate correlating {r}: power for 0.30 SD at n = {n} is {pw:.3f}")
# standardised effect 0.20:  394 per arm for 80 % power,  527 for 90 %
# standardised effect 0.25:  253 per arm for 80 % power,  338 for 90 %
# standardised effect 0.30:  176 per arm for 80 % power,  235 for 90 %
# standardised effect 0.40:  100 per arm for 80 % power,  133 for 90 %
# chosen: 0.30 SD, 90 % power -> 235 analysable per arm
# recruit 277 per arm allowing 15 % exclusions (comprehension checks), 554 in all
# with a baseline covariate correlating 0.5: power for 0.30 SD at n = 235 is 0.964