#!/usr/bin/python
# vim: set fileencoding=utf-8 :

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

import datetime
import re
import requests as req
import locale
import os.path

locale.setlocale(locale.LC_ALL, 'de_DE.UTF-8')

sources_folder = 'Quellen'
plots_folder = 'Plots'


einwohner_deutschland = 83190556

today = datetime.date.today()
print_today = today.isoformat()

filename_now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

# DIN A4 Plots
plt.rcParams["figure.figsize"] = [11.69, 8.27]


# Download

filename = '{}/{}_Impfquotenmonitoring.xlsx'.format(sources_folder, filename_now)

r = req.get('https://www.rki.de/DE/Content/InfAZ/N/Neuartiges_Coronavirus/Daten/Impfquotenmonitoring.xlsx?__blob=publicationFile')

with open(filename, 'wb') as outfile:
	outfile.write(r.content)




rki_file = pd.read_excel(filename, sheet_name=None, engine='openpyxl')

raw_data = rki_file['Impfungen_proTag']

impfungen = raw_data[:-1].dropna()

dates = impfungen['Datum']
daily = impfungen['Gesamtzahl Impfungen']
cumulative = np.cumsum(impfungen['Gesamtzahl Impfungen'])

total_vaccinations = int(np.sum(daily))
total_vaccinations_percentage = float(total_vaccinations) / einwohner_deutschland

mean_vaccinations_daily = np.mean(daily)
mean_vaccinations_daily_int = int(np.round(mean_vaccinations_daily))

to_be_vaccinated = einwohner_deutschland - total_vaccinations
days_extrapolated = int(np.ceil(to_be_vaccinated / mean_vaccinations_daily))

extrapolated_dates = np.array([dates[0] + datetime.timedelta(days=i) for i in range(days_extrapolated)])

extrapolated_vaccinations = mean_vaccinations_daily * range(days_extrapolated)

mean_vaccinations_daily_up_to_date  = np.round(cumulative / range(1, len(cumulative) + 1))

# Stand aus Daten auslesen
#stand = dates.iloc[-1]
#print_stand = stand.isoformat()

# Stand aus offiziellen Angaben auslesen
stand = rki_file['Erläuterung'].iloc[1][0]

stand_regex = re.compile('^Datenstand: (\d\d.\d\d.\d\d\d\d, \d\d:\d\d) Uhr$')
m = stand_regex.match(stand)
stand_date = datetime.datetime.strptime(m.groups()[0], '%d.%m.%Y, %H:%M')
print_stand = stand_date.isoformat()

filename_stand = stand_date.strftime("%Y%m%d%H%M%S")

def plot_extrapolation_portion(percentage):

	print_percentage = int(percentage * 100)
	plot_filename = '{}/{}_extrapolated_to_{}_percent.pdf'.format(plots_folder, filename_stand, print_percentage)
	plot_filename_latest = '{}/latest_extrapolated_to_{}_percent.pdf'.format(plots_folder, print_percentage)

	if os.path.isfile(plot_filename):
		print('Plot {} already exists'.format(plot_filename))
		return

	fig, ax = plt.subplots(1)


	plt.title(
		'Tägliche Impfquote, kumulierte Impfungen und lineare Extrapolation bis {:n} % der Bevölkerung Deutschlands\n'
		'Erstellung: {}, Datenquelle: RKI, Stand: {}\n'
		'Impfungen gesamt: {:n} ({:n} %), Durchschnittliche Impfrate: {:n} Impfungen/Tag'.format(
			print_percentage,
			print_today, print_stand,
			total_vaccinations, np.round(total_vaccinations_percentage * 100, 2), mean_vaccinations_daily_int
		)
	)

	ax2 = ax.twinx()

	ax.bar(dates, daily, label='Tägliche Impfungen', color='blue')
	ax.plot(dates, mean_vaccinations_daily_up_to_date, color='violet', label='Durchschnittliche Impfquote\nbis zu diesem Tag (inkl.)')

	ax2.set_ylim([0, einwohner_deutschland * percentage])
	ax2.set_xlim(xmax=dates[0] + datetime.timedelta(days=percentage * days_extrapolated))
	ax2.grid(True)
	ax2.plot(dates, cumulative, color='red', label='Kumulierte Impfungen')
	ax2.plot(extrapolated_dates, extrapolated_vaccinations, color='orange', label='Extrap. kumulierte Impfungen\n({:n} Impfungen/Tag)'.format(mean_vaccinations_daily_int))
	#ax2.plot()

	ax.legend(loc='upper left')
	ax.get_yaxis().get_major_formatter().set_scientific(False)

	ax.set_xlabel('Datum')
	ax.set_ylabel('Tägliche Impfungen')

	ax2.legend(loc='center right')
	ax2.get_yaxis().get_major_formatter().set_scientific(False)

	# Estimated percentage for herd immunity
	#ax2.axline((0, einwohner_deutschland * 0.7), slope=0, color='green')

	ax2.set_ylabel('Kumulierte Impfungen')

	plt.savefig(plot_filename)
	plt.savefig(plot_filename_latest)
	plt.close()

	print('Created plot {}'.format(plot_filename))


plot_extrapolation_portion(0.1)
plot_extrapolation_portion(1.0)