"""Download alongside portfolio-results.json, then run with Python 3.
Uses only the standard library. All percentages are converted to decimal units.
"""
from pathlib import Path
import json
import math

data = json.loads(Path(__file__).with_name('portfolio-results.json').read_text())
assets = {a['id']: a for a in data['inputs']['assets']}
indices = {asset: i for i, asset in enumerate(assets)}

def calculate(weights, correlation_stress=False, return_change=0):
    def correlation(a, b):
        value = data['inputs']['corr'][indices[a]][indices[b]]
        return 0.5 * value + 0.5 if correlation_stress else value

    variance = sum(
        weights[a] * weights[b] * assets[a]['v'] * assets[b]['v']
        / 10000 * correlation(a, b)
        for a in weights for b in weights
    )
    arithmetic = sum(
        weight * (assets[a]['r'] / 100 + return_change + (assets[a]['v'] / 100) ** 2 / 2)
        for a, weight in weights.items()
    )
    return 100 * (arithmetic - variance / 2), 100 * math.sqrt(variance)

for name, row in data['portfolios'].items():
    weights = {a: value / 100 for a, value in row['weights_pct'].items()}
    result = calculate(weights)
    assert abs(result[0] - row['compound_return_pct']) < 1e-10
    assert abs(result[1] - row['volatility_pct']) < 1e-10
    print(f'{name}: return {result[0]:.2f}%, volatility {result[1]:.2f}%')

weights = {a: value / 100 for a, value in data['portfolios']['Strategic illustration']['weights_pct'].items()}
for label, stress, change in [('strategic_correlations_halfway_to_one', True, 0), ('strategic_all_asset_returns_minus_2pp', False, -0.02)]:
    result = calculate(weights, stress, change)
    expected = data['sensitivity'][label]
    assert abs(result[0] - expected['compound_return_pct']) < 1e-10
    assert abs(result[1] - expected['volatility_pct']) < 1e-10
    print(f'{label}: return {result[0]:.2f}%, volatility {result[1]:.2f}%')