import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

stock_data = pd.read_csv('A.csv',parse_dates=['Datetime'])
stock_data['Changes_Next_Year'] = stock_data.groupby('Code')['Change'].shift(-1)
stock_data.dropna(subset=['Changes_Next_Year'],inplace=False)

def simulate(stock_data, NumberOfShares):
    '''
    Randomly select n number of shares
    and check their annual performance
    '''
    stock_data = stock_data.groupby(['Datetime']).sample(NumberOfShares)
    output = pd.DataFrame()
    stock_data['Code'] += ' '
    stock_data_groupby = stock_data.groupby('Datetime')
    output['Shares_bought'] = stock_data_groupby['Code'].sum()
    output['Number of shares'] = stock_data_groupby.size()
    output['Average gains'] = stock_data_groupby['Changes_Next_Year'].mean()
    output['Average gains'].replace(np.inf, 0, inplace=True)

    output['Capital'] = 0
    output = output.dropna()

    output['Capital'] = (output['Average gains']+1.0).cumprod()
    output['Average gains'] = output['Average gains']*100.0
    
    return output['Average gains'].mean(), output['Average gains'].std()

print(simulate(stock_data, 2))

def run(stock_data,NumberOfShares):
    ret = 0
    std = 0

    iteration = 60
    for i in range(1,iteration + 1):
        r,s = simulate(stock_data,NumberOfShares)
        ret += r
        std += s

    return ret/(iteration), std/(iteration)

# simulate from 1 stock to 100 stocks portfolio
stdev = []
ret = []
for i in range(1,100):
    tempr, temps = run(stock_data,i)
    stdev.append(temps)
    ret.append(tempr)

# visualize the data
fig, ax1 = plt.subplots(figsize=(10,5))
plt.style.use('seaborn')

ax2 = ax1.twinx()
ax1.plot(stdev, 'darkgreen',linestyle = '-', marker = '.', alpha = 0.5)
ax2.plot(ret, 'darkblue', linestyle = '-', marker = '.', alpha = 0.5)

ax1.set_xlabel('Number of Shares in portfolio')
ax1.set_ylabel('Standard Deviation of Returns', color='darkgreen')
ax2.set_ylabel('Average Compounded Annual Growth Rate', color='darkblue')
ax1.grid(False)
ax2.grid(False)

plt.show()