"""Reproduce the hypothetical examples in D1, 7 September 2026.

Python 3.10+, standard library only. Run:
    python certainty_equivalent_calculations.py

No inputs are calibrated forecasts. Returns follow geometric Brownian motion
until release and a known one-year return thereafter. All rates are continuously
compounded where expressed as discount rates. No fees, tax or inflation.
"""

import json
import math
from statistics import NormalDist


MU = 0.06           # GBM proportional drift; E[G_t] = exp(MU * t).
SIGMA = 0.10        # Annual diffusion coefficient.
GAMMA = 3.0        # CRRA relative risk aversion.
POST_RELEASE = 0.04 # Known, continuously compounded rate to payment.


def investment_factors(release, payment):
    """Investment only: no mortality credits, no own-survival conditioning."""
    mean_log = (MU - 0.5 * SIGMA ** 2) * release
    mean_log += POST_RELEASE * (payment - release)
    variance_log = SIGMA ** 2 * release
    mean_factor = math.exp(mean_log + 0.5 * variance_log)
    ce_factor = math.exp(mean_log + 0.5 * (1 - GAMMA) * variance_log)
    return mean_log, variance_log, mean_factor, ce_factor


def normal_integral(function, bound=10.0, steps=20000):
    """Simpson quadrature over [-10,10]; omitted tails are negligible here."""
    h = 2 * bound / steps
    normalizer = math.sqrt(2 * math.pi)

    def integrand(z):
        return function(z) * math.exp(-z * z / 2) / normalizer

    total = integrand(-bound) + integrand(bound)
    for i in range(1, steps):
        total += (4 if i % 2 else 2) * integrand(-bound + i * h)
    return total * h / 3


def mortality_ce(n, survival):
    """Toy homogeneous terminal pool, conditional on one member surviving."""
    expected_power = 0.0
    for other_survivors in range(n):
        probability = math.comb(n - 1, other_survivors)
        probability *= survival ** other_survivors
        probability *= (1 - survival) ** (n - 1 - other_survivors)
        uplift = n / (1 + other_survivors)
        expected_power += probability * uplift ** (1 - GAMMA)
    enumerated = expected_power ** (1 / (1 - GAMMA))
    # The closed form below specifically uses GAMMA=3.
    mean_survivors = 1 + (n - 1) * survival
    variance_survivors = (n - 1) * survival * (1 - survival)
    closed_form = n / math.sqrt(mean_survivors ** 2 + variance_survivors)
    assert math.isclose(enumerated, closed_form, rel_tol=1e-12)
    return enumerated


def main():
    target, release, payment = 10000.0, 19.0, 20.0
    survival = 0.80
    m, v, mean_factor, ce_factor = investment_factors(release, payment)

    # Independent numerical integration of E[G^(1-gamma)], before inversion.
    expected_power = normal_integral(
        lambda z: math.exp((1 - GAMMA) * (m + math.sqrt(v) * z))
    )
    numerical_ce = expected_power ** (1 / (1 - GAMMA))
    assert math.isclose(numerical_ce, ce_factor, rel_tol=1e-10)

    normal = NormalDist()
    ce_capital = target / ce_factor
    # Under CE funding, log(payment/target) has the following normal mean.
    relative_mean_log = (GAMMA - 1) * v / 2
    q05_factor = math.exp(relative_mean_log + normal.inv_cdf(.05) * math.sqrt(v))

    rungs = []
    for payment_year in range(2, 6):
        release_year = payment_year - 1
        _, _, _, factor = investment_factors(release_year, payment_year)
        rungs.append({
            "release_year": release_year,
            "payment_year": payment_year,
            "initial_investment": 1000.0 / factor,
            "nominal_target": 1000.0,
        })

    finite_pools = []
    for n in (5, 100):
        uplift_ce = mortality_ce(n, survival)
        finite_pools.append({
            "original_members": n,
            "release_survival_probability": survival,
            "mortality_ce_uplift": uplift_ce,
            "ideal_deterministic_uplift": 1 / survival,
            "additional_capital_relative_to_ideal": 1 / (survival * uplift_ce) - 1,
            "combined_ce_capital_under_independence": ce_capital / uplift_ce,
        })

    result = {
        "assumptions": {
            "gbm_drift": MU, "annual_volatility": SIGMA,
            "risk_aversion_gamma": GAMMA,
            "post_release_continuous_yield": POST_RELEASE,
            "fees_tax_inflation": "excluded; nominal amounts only",
            "mortality_investment_independence": True,
        },
        "four_rung_investment_only_ladder": rungs,
        "four_rung_initial_total": sum(r["initial_investment"] for r in rungs),
        "longer_rung": {
            "target": target, "release_year": release, "payment_year": payment,
            "mean_log_return": m, "variance_log_return": v,
            "mean_gross_growth": mean_factor, "ce_gross_growth": ce_factor,
            "ce_from_numerical_integration": numerical_ce,
            "annualised_ce_continuous_rate": math.log(ce_factor) / payment,
            "annualised_ce_effective_rate": ce_factor ** (1 / payment) - 1,
            "mean_funding": target / mean_factor,
            "ce_funding": ce_capital,
            "matched_bond_funding_at_hypothetical_4pct_continuous": target * math.exp(-.04 * payment),
            "ideal_survival_adjusted_ce_funding": survival * ce_capital,
            "ideal_survival_adjusted_mean_funding": survival * target / mean_factor,
            "ce_funded_probability_below_target": normal.cdf(-relative_mean_log / math.sqrt(v)),
            "ce_funded_median_payout": target * math.exp(relative_mean_log),
            "ce_funded_fifth_percentile_payout": target * q05_factor,
            "mean_funded_probability_below_target": normal.cdf(math.sqrt(v) / 2),
            "capital_for_95pct_target_threshold_same_risky_strategy": ce_capital / q05_factor,
        },
        "toy_finite_pools_not_FTP": finite_pools,
        "checks": "CE quadrature and finite-pool binomial sums agree with closed forms",
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()