#!/usr/bin/python

import math
import numpy as np

num_stocks = 100000
mean_initial_fair_price = 100.0
stddev_initial_fair_price = 10.0
stddev_noise = 0.01
mean_num_shares = 1000
stddev_num_shares = 100
# 1.0004^250 is approximately 1.1.  In other words, if 1.0004 is the
# mean return for a day, then 1.1 (a ten percent gain) is the mean
# return for a year (there are roughly 250 trading days in a year).
mean_return = 1.0004
stddev_return = 0.01
num_days = 250
sum_cum_market_cap_returns = 0
sum_day_market_cap_returns = 0
sum_cum_ew_returns = 0
sum_day_ew_returns = 0

def day():
    global fair_prices, market_prices
    market_caps = np.multiply(market_prices, num_shares)
    total_market_cap = np.sum(market_caps)
    cap_weights = np.divide(market_caps, total_market_cap)
    returns = np.random.normal(mean_return, stddev_return, num_stocks)
    fair_prices = np.multiply(fair_prices, returns)
    noise = np.random.normal(1.0, stddev_noise, num_stocks)
    market_prices = np.multiply(fair_prices, noise)
    cap_weighted_return = np.dot(cap_weights, returns)
    equally_weighted_return = np.sum(returns) / num_stocks
    return (cap_weighted_return, equally_weighted_return)

fair_prices = np.random.normal(mean_initial_fair_price,
                               stddev_initial_fair_price, num_stocks)
noise = np.random.normal(1.0, stddev_noise, num_stocks)
market_prices = np.multiply(fair_prices, noise)
num_shares = np.random.normal(mean_num_shares, stddev_num_shares,
                              num_stocks)
cum_cap_weighted_return = 1.0
cum_equally_weighted_return = 1.0
for d in range(num_days):
    (cap_weighted_return, equally_weighted_return) = day()
    cum_cap_weighted_return *= cap_weighted_return
    cum_equally_weighted_return *= equally_weighted_return

print('Cum cap weighted return %f' % cum_cap_weighted_return)
print('Cum equally weighted return %f' % cum_equally_weighted_return)
daily_cap_weighted_return = \
math.pow(cum_cap_weighted_return, 1.0 / num_days)
daily_equally_weighted_return = \
math.pow(cum_equally_weighted_return, 1.0 / num_days)
print('Daily cap weighted return %f' % daily_cap_weighted_return)
print('Daily equally weighted return %f' % daily_equally_weighted_return)