#!/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, 4, 6))
    
    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
        return v / max(v)

# 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')

ax.xaxis.set_major_locator(mpl.dates.MonthLocator())
ax.xaxis.set_major_formatter(mpl.dates.DateFormatter("%Y-%m"))
ax.yaxis.set_major_formatter(mpl.ticker.FormatStrFormatter('%.0f%%'))

plt.axvline(datetime.datetime(2020, 3, 9), color='k')
ax.text(datetime.datetime(2020, 3, 9), 0.64, 'BMI', fontsize=11, ha='left', va='bottom',
    transform=mpl.transforms.blended_transform_factory(ax.transData, ax.transAxes))
ax.text(datetime.datetime(2020, 3, 16), 0.51, 'BMII', fontsize=11, ha='left', va='bottom',
    transform=mpl.transforms.blended_transform_factory(ax.transData, ax.transAxes))
plt.axvline(datetime.datetime(2020, 3, 16), color='k')

plt.xlabel('date')
plt.ylabel('value relative to 2020 maximum')
plt.grid(True)
plt.legend(loc='center left', framealpha=1, edgecolor='k', borderpad=0.7, borderaxespad=2)
plt.tight_layout()
plt.savefig('stock-indices-2020crash.svg')
plt.show()