import math
def trade(stock, length):
mu = pd.rolling_mean(stock, window=length)
portfolio=[]
portfolio.append(stock[0])
#for the initial section, just assume does what SPY does
current_shares=1
for i in range(1,len(stock)):
#stock is less than rolling average, sell at end of day
if stock[i]<mu[i] and current_shares>0:
#end of day portfolio will be however much realized from sale
portfolio.append(current_shares*stock[i])
current_shares=0
# if stock > mu, buy, if stock<mu, sell
elif stock[i]>=mu[i] and current_shares==0:
#buy with however much money had on prior day
current_shares=portfolio[i-1]/stock[i]
portfolio.append(current_shares*stock[i])
elif stock[i]<mu[i] and current_shares==0:
portfolio.append(portfolio[i-1])
elif stock[i]>=mu[i] and current_shares>0:
portfolio.append(current_shares*stock[i])
#for rolling index is NaN, etc.
elif math.isnan(mu[i]):
portfolio.append(current_shares*stock[i])
return portfolio, mu