#!/usr/bin/env python3
"""Exact synthetic counterexample; not a reproduction of a published model.

Five paths have explicit probabilities and wealth at t=0, 1, 2 in USD.
Keep the two most probable paths and condition on that retained set.
All calculations use fractions; no random seed or third-party package is needed.
"""

from fractions import Fraction as F


SCENARIOS = (
    (F(47, 100), (100, 99, 99)),
    (F(47, 100), (100, 101, 101)),
    (F(4, 100), (100, 70, 100)),
    (F(1, 100), (100, 60, 60)),
    (F(1, 100), (100, 140, 140)),
)


def quantile(distribution, level):
    """Left quantile: smallest x whose cumulative probability is >= level."""
    cumulative = F(0)
    for value, probability in sorted(distribution):
        cumulative += probability
        if cumulative >= level:
            return value
    raise ValueError("Distribution does not reach the requested level")


def expected_shortfall(losses, confidence=F(95, 100)):
    """Average the worst 1-confidence probability mass, splitting atoms."""
    tail_mass = F(1) - confidence
    remaining, total = tail_mass, F(0)
    for loss, probability in sorted(losses, reverse=True):
        included = min(probability, remaining)
        total += included * loss
        remaining -= included
        if remaining == 0:
            return total / tail_mass
    raise ValueError("Not enough probability mass for the tail")


def maximum_drawdown(path):
    peak, drawdown = F(path[0]), F(0)
    for wealth in path:
        peak = max(peak, F(wealth))
        drawdown = max(drawdown, F(1) - F(wealth) / peak)
    return drawdown


def summarize(scenarios):
    assert sum(probability for probability, _ in scenarios) == 1
    wealth = [(F(path[-1]), probability) for probability, path in scenarios]
    losses = [(F(1) - F(path[-1], path[0]), probability)
              for probability, path in scenarios]
    drawdowns = [(maximum_drawdown(path), probability)
                 for probability, path in scenarios]
    return (
        sum(value * probability for value, probability in wealth),
        quantile(wealth, F(5, 100)),
        quantile(losses, F(95, 100)),
        expected_shortfall(losses),
        sum(value * probability for value, probability in drawdowns),
        sum(probability for value, probability in drawdowns if value >= F(1, 5)),
    )


def main():
    retained = sorted(SCENARIOS, key=lambda item: item[0], reverse=True)[:2]
    retained_mass = sum(probability for probability, _ in retained)
    truncated = tuple((probability / retained_mass, path)
                      for probability, path in retained)
    full_stats, truncated_stats = summarize(SCENARIOS), summarize(truncated)
    assert retained_mass == F(94, 100)
    assert full_stats == (F(100), F(99), F(1, 100), F(88, 1000),
                          F(207, 10000), F(5, 100))
    assert truncated_stats == (F(100), F(99), F(1, 100), F(1, 100),
                               F(5, 1000), F(0))
    print("Synthetic exact-probability example; initial wealth = 100 USD")
    print(f"Retained probability mass: {float(retained_mass):.2%}")
    print("law       mean_USD q05_USD VaR95_loss ES95_loss mean_MDD P(MDD>=20%)")
    for name, stats in (("full", full_stats), ("truncated", truncated_stats)):
        mean, q05, var, es, mdd, breach = map(float, stats)
        print(f"{name:<9} {mean:8.2f} {q05:7.2f} {var:10.2%} "
              f"{es:9.2%} {mdd:8.2%} {breach:11.2%}")
    print("Exact arithmetic assertions: PASS")


if __name__ == "__main__":
    main()
