#!/usr/bin/python3

# Investing Math Model - Simple model of unleveraged and leveraged investment return probability density functions.
# Copyright (C) 2017 Gordon Irlam
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from functools import reduce
from math import ceil, exp, floor, log, sqrt
from operator import mul
from statistics import mean

from scipy.stats import lognorm

asset_classes = (
    {
        # Credit-Suiss Yearbook 2017: Global 1900-2016.
        'name': 'global-historical',
        'm': 1.065,
        'vol': 0.174,
        'leverage': 1,
        'borrow': 0,
        'expense': 0,
        'r_lo': -0.1,
        'r_hi': 0.2,
    },
    # {
    #     # Shiller: U.S. Large Cap 1927-2015.
    #     'name': 'us_large_historical',
    #     'm': 1.086,
    #     'vol': 0.195,
    # },
    # {
    #     # Fama and French: U.S. Small Cap Value 1927-2015.
    #     'name': 'us_small_value_historical',
    #     'm': 1.157,
    #     'vol': 0.314,
    # },
    {
        # Projected: U.S. Large Cap.
        'name': 'investing-us_large',
        'm': 1.045,
        'vol': 0.168,
        'leverage': 1,
        'borrow': 0,
        'expense': 0,
    },
    {
        # Projected: U.S. Small Cap Value.
        'name': 'investing-us_small_value',
        'm': 1.065,
        'vol': 0.222,
        'leverage': 1,
        'borrow': 0,
        'expense': 0,
    },
    {
        # Projected: 3X U.S. Large Cap. margin current borrowing cost
        'name': 'investing-3x_us_large_current',
        'm': 1.045,
        'vol': 0.168,
        'leverage': 3,
        'borrow': -0.001,
            # June 16, 2017: Interactive Brokers charges 2.16% nominal 1% above FEDFUNDS benchmark rate for $100k-$1m.
            # Survey of Professional Forecastors: Inflation rate expectation on 2017Q2 for 2017 2.3%.
            # 2.16 - 2.3 = -0.1%.
        'expense': 0,
    },
    {
        # Projected: 3X U.S. Large Cap. margin historical borrowing cost
        'name': 'investing-3x_us_large_historical',
        'm': 1.045,
        'vol': 0.168,
        'leverage': 3,
        'borrow': 0.022,
            # June 16, 2017: Interactive Brokers charges 1% nominal above FEDFUNDS benchmark rate for $100k-$1m.
            # June 16, 2017, FEDFUNDS exceeds 1-month cmt by 0.31%.
            # Historical risk free 1900-2000 was 0.9% real.
            # 0.9 + 1 + 0.31 = 2.2%.
        'expense': 0,
    },
    {
        # Projected: 3X U.S. Large Cap. ETF current borrowing cost
        'name': 'leveraged_etfs-3x_us_large_current',
        'm': 1.045,
        'vol': 0.168,
        'leverage': 3,
        'borrow': -0.014,
            # 2017-04 Fred FEDFUNDS: 0.90%.
            # 2017Q2 SPF: Inflation rate 2017: 2.3%.
            # 0.90 - 2.3 = -1.4%.
        'expense': 0.018,
            # 2010-2016 UPRO NAV.
    },
    {
        # Projected: 3X U.S. Large Cap. ETF historical borrowing cost
        'name': 'leveraged_etfs-3x_us_large_historical',
        'm': 1.065,
        'vol': 0.174,
        'leverage': 3,
        'borrow': 0.012,
            # No Interactive Brokers 1% premium.
        'expense': 0.018,
            # 2010-2016 UPRO NAV.
    },
    {
        # Projected: 2X U.S. Large Cap. ETF current borrowing cost
        'name': 'leveraged_etfs-2x_us_large_current',
        'm': 1.045,
        'vol': 0.168,
        'leverage': 2,
        'borrow': -0.014,
            # 2017-04 Fred FEDFUNDS: 0.90%.
            # 2017Q2 SPF: Inflation rate 2017: 2.3%.
            # 0.90 - 2.3 = -1.4%.
        'expense': 0.014,
            # 2007-2016 SSO.
    },
    {
        # Projected: 2X U.S. Large Cap. ETF historical borrowing cost
        'name': 'leveraged_etfs-2x_us_large_historical',
        'm': 1.065,
        'vol': 0.174,
        'leverage': 2,
        'borrow': 0.012,
        'expense': 0.014,
            # 2007-2016 SSO.
    },
)

for ac in asset_classes:

    m = ac['m']
    vol = ac['vol']
    #frequency = ac['frequency']
        # Frequency is by assumption continuous.

    base_sigma = sqrt(log(1 + (vol / m) ** 2))
    base_mu = log(m / sqrt(1 + (vol / m) ** 2))
    base_mu_gbm = base_mu + base_sigma ** 2 / 2 # Or equivalently, log(m).
    mu_gbm = ac['leverage'] * base_mu_gbm + (1 - ac['leverage']) * log(1 + ac['borrow']) + log(1 - ac['expense'])
    sigma = ac['leverage'] * base_sigma
    mu = mu_gbm - sigma ** 2 / 2

    mn = exp(mu + sigma ** 2 / 2)
    vl = sqrt((exp(sigma ** 2) - 1) * exp(2 * mu + sigma ** 2))

    print('%s, leverge %f, borrow %f, expense %f:' % (ac['name'], ac['leverage'], ac['borrow'], ac['expense']))
    print('    mean %f, vol %f, gm %f' % (mn - 1, vl, exp(mu) - 1))
    print('    mu %f, sigma %f' % (mu, sigma))

    # n_high = 50
    # n_steps = 2

    # with open('investing-' + ac['name'] + '-' + str(leverage) + '-risk_percentiles.csv', 'w') as f:
    #     pctls = (0.01, 0.05, 0.2, 0.5)
    #     for i in range(1, n_high * n_steps + 1):
    #         n = i / n_steps
    #         test_slow = False
    #         if test_slow and i % n_steps == 0:
    #             # Slow, approximate method (useful to verify fast, exact formulas):
    #             samples = 10000
    #             returns = sorted(reduce(mul, lognorm.rvs(sigma, scale=exp(mu), size=n), 1) ** (1/n) - 1 for _ in range(samples))
    #             pctl_returns = (returns[int(pctl * samples + 0.5)] for pctl in pctls)
    #             print(n, mean(returns), *pctl_returns)
    #         # Fast, exact method:
    #         pctl_returns = (lognorm.ppf(pctl, sigma/sqrt(n), scale=exp(mu)) - 1 for pctl in pctls)
    #         f.write('%f,%f' % (n, lognorm.stats(sigma/sqrt(n), scale=exp(mu), moments='m') - 1))
    #         for p in pctl_returns:
    #             f.write(',%f' % p)
    #         f.write('\n')

    years = (10, 20, 50)

    for n in years:
        print('    %d years, median factor %f, mean factor %f' %
            (n, lognorm.median(sqrt(n) * sigma, scale = exp(n * mu)),  lognorm.mean(sqrt(n) * sigma, scale = exp(n * mu))))

    f_max = 10
    f_steps = 100

    step = f_max / f_steps
    with open(ac['name'] + '-final_value_factor_pdf.csv', 'w') as f, open(ac['name'] + '-final_value_factor_cdf.csv', 'w') as g:
        for i in range(ceil(f_max * f_steps)):
            fvf = i / f_steps
            pd = tuple(lognorm.pdf(fvf, sqrt(n) * sigma, scale = exp(n * mu)) for n in years)
            cd = tuple(lognorm.cdf(fvf, sqrt(n) * sigma, scale = exp(n * mu)) for n in years)
            f.write('%f' % fvf)
            g.write('%f' % fvf)
            for p in pd:
                f.write(',%f' % p)
            for c in cd:
                g.write(',%f' % c)
            f.write('\n')
            g.write('\n')

    r_lo = ac.get('r_lo', -0.20)
    r_hi = ac.get('r_hi', 0.40)
    r_steps = 1000

    with open(ac['name'] + '-annualized_roi_pdf.csv', 'w') as f, open(ac['name'] + '-annualized_roi_cdf.csv', 'w') as g:
        for i in range(floor(r_lo * r_steps), ceil(r_hi * r_steps)):
            r = i / r_steps
            pd = tuple(lognorm.pdf(1 + r, sigma / sqrt(n), scale = exp(mu)) for n in years)
            cd = tuple(lognorm.cdf(1 + r, sigma / sqrt(n), scale = exp(mu)) for n in years)
            f.write('%f' % r)
            g.write('%f' % r)
            for p in pd:
                f.write(',%f' % p)
            for c in cd:
                g.write(',%f' % c)
            f.write('\n')
            g.write('\n')

    print('    0.025 ppfs', dict({n: lognorm.ppf(0.025, sigma / sqrt(n), scale = exp(mu)) - 1 for n in years}))
