# import pandas, numpy, and matplotlib
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
pd.set_option('display.width', 80)
pd.set_option('display.max_columns', 7)
pd.set_option('display.max_rows', 200)
pd.options.display.float_format = '{:,.2f}'.format
covidtotals = pd.read_csv("dane/covidtotals.csv", parse_dates=["lastdate"])

# generate a correlation matrix

corr = covidtotals.corr()
corr[['total_cases','total_deaths','total_cases_pm','total_deaths_pm']]

# show scatter plots
fig, axes = plt.subplots(1,2, sharey=True)
sns.regplot(covidtotals.median_age, covidtotals.total_cases_pm, ax=axes[0])
sns.regplot(covidtotals.gdp_per_capita, covidtotals.total_cases_pm, ax=axes[1])
axes[0].set_xlabel("Wiek")
axes[0].set_ylabel("Liczba przypadków na milion")
axes[1].set_xlabel("PKB per capita")
axes[1].set_ylabel("")
plt.suptitle("Wykres punktowy liczby przypadków na milion w zależności od wieku i PKB per capita")
plt.tight_layout()
fig.subplots_adjust(top=0.92)
plt.show()

# generate a heat map
sns.heatmap(corr, xticklabels=corr.columns, yticklabels=corr.columns, cmap="coolwarm")
plt.title('Mapa ciepła macierzy korelacji')
plt.tight_layout()
plt.show()

