#!/usr/bin/python3
# -*- coding: utf8 -*-

import csv
import datetime
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

class Stock:
    def __init__(self, name):
        self.data = self.get_csv(name)
        self.convert_types()
        self.filter_date(datetime.datetime(2020, 1, 1), datetime.datetime(2020, 12, 31))
    
    def get_csv(self, name):
        try:
            with open(name, 'r' ) as f:
                reader = csv.DictReader(f)
                return [line for line in reader]
        except FileNotFoundError as ex:
            print(ex)
            print('get data from', 'https://finance.yahoo.com/quote/DAX/history?p=DAX')
            exit()
    
    def convert_types(self):
        for il, l in enumerate(self.data):
            for k in l.keys():
                try:
                    if k == 'Date':
                        l[k] = datetime.datetime.strptime(l[k], '%Y-%m-%d')
                    else:
                        l[k] = float(l[k])
                except Exception:
                    del self.data[il]
    
    def filter_date(self, date, date2=None):
        self.data = [i for i in self.data if i['Date'] >= date]
        if date2 is not None:
            self.data = [i for i in self.data if i['Date'] <= date2]
    
    def get_dates(self):
        return [l['Date'] for l in self.data]
    
    def get_values(self):
        return np.array([float(l['Close']) for l in self.data])
    
    def get_values_norm(self):
        v = self.get_values()
        #vmean = np.mean([v for i, v in enumerate(v) if self.data[i]['Date'].month == 1])
        #return v / vmean
        # v / max(v)
        return v / max([v for i, v in enumerate(v) if self.data[i]['Date'].month < 4])


# data is found on finance.yahoo.com
data_spx = Stock('^GSPC.csv')
data_DJI = Stock('^DJI.csv')
data_stoxx50e = Stock('^STOXX50E.csv')
data_DAX = Stock('^GDAXI.csv')

plt.figure(figsize=[5.6, 4.2])
ax = plt.gca()
ax.set_prop_cycle(color=['#0072bd', '#d95319', '#edb120', '#7e2f8e'])

plt.plot(data_spx.get_dates(), 100*data_spx.get_values_norm(), 'o-', ms=3, label='S&P 500')
plt.plot(data_DJI.get_dates(), 100*data_DJI.get_values_norm(), 'o-', ms=3, label='Dow Jones')
plt.plot(data_stoxx50e.get_dates(), 100*data_stoxx50e.get_values_norm(), 'o-', ms=3, label='EURO STOXX 50')
plt.plot(data_DAX.get_dates(), 100*data_DAX.get_values_norm(), 'o-', ms=3, label='DAX')

plt.xlim(datetime.datetime.strptime('2019-12-24', '%Y-%m-%d'),
         datetime.datetime.strptime('2021-01-08', '%Y-%m-%d'))
ax.tick_params(axis='x', which='minor', pad=6)
ax.xaxis.set_major_locator(mpl.dates.MonthLocator())
ax.xaxis.set_major_formatter(mpl.ticker.NullFormatter())
ax.xaxis.set_minor_formatter(mpl.dates.DateFormatter("%m"))
ax.xaxis.set_minor_locator(mpl.dates.MonthLocator(bymonthday=16))
ax.yaxis.set_major_formatter(mpl.ticker.FormatStrFormatter('%.0f%%'))

for tick in ax.xaxis.get_minor_ticks():
    tick.tick1line.set_markersize(0)
    tick.tick2line.set_markersize(0)
    tick.label1.set_horizontalalignment('center')

plt.xlabel('date (month 2020)')
plt.ylabel('value relative to 2020 pre-crash maximum')
plt.grid(True)
plt.legend(loc='lower right', framealpha=1, edgecolor='k', borderpad=0.7, borderaxespad=0.6)
plt.tight_layout()
plt.savefig('stock-indices-2020crash+recovery.svg')