import csv
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from datetime import datetime

# period (days) for consideration
P = 21 

# read in the GEX+ Spreadsheet
with open('GEX.csv') as f:
    data = [row for row in csv.DictReader(f)]

# make a parallel list of P-period % returns
returns = []
for i, d in enumerate(data):
    if i < len(data) - P:
        now = float(data[i]['CLOSE'])
        then = float(data[i+P]['CLOSE'])
        chg = (then - now) / now
    else:
        chg = None

    returns.append(chg)

# 
events = []	# list indices of days that fit the criteria
histogram = []	# returns of days that fit the criteria
for i, d in enumerate(data):
    cr = float(d['CR(x)'])	# Crash Risk
    npd = float(d['NPD'])	# Net Put Delta (NPD)
    gex = float(d['GEX+'])
    vgr = float(d['VGR'])
    if 5 > npd > -5 and vgr > 0:
        histogram.append(returns[i])
        events.append(i)

# clean up histogram data to remove extraneous None-types
histogram = list(filter(None, histogram))

# isolate SPX closes for plotting
timeseries = [float(d['CLOSE']) for d in data]
timeseries_x = [n for n in range(0, len(timeseries))]

# create date formatter
dates = [datetime.strptime(d['DATE'], '%Y-%m-%d') for d in data]
boy = []	# beginning of year markers
for i, n in enumerate(data):
    if i > 0 and dates[i].year != dates[i-1].year:
        boy.append(i)

def format_date(idx, pos=None):
    idx = min(int(idx), len(dates) - 1)
    return dates[idx].strftime("'%y")

# draw plot
fig = plt.figure(figsize=(12, 6))
left, bottom, width, height = 0.1, 0.1, 0.8, 0.40
ax1 = fig.add_axes( (left, bottom, width, height) )
ax1.hist(histogram, max(10, int(len(histogram)/10)))		# 10 bins

left, bottom, width, height = 0.1, 0.55, 0.8, 0.40
ax2 = fig.add_axes( (left, bottom, width, height) )
ax2.plot(timeseries_x, timeseries)
ax2.scatter(events, [timeseries[n]+250 for n in events], c='red', lw=0, s=10)
ax2.set_xlim([0, max(timeseries_x)+100])
ax2.xaxis.set_major_formatter(ticker.FuncFormatter(format_date))
ax2.set_xticks(boy)

print('EVENTS:            ', len(histogram))
print('MEAN %s-day return:' % P, sum(histogram) / len(histogram))

plt.show()