# Visualization of rate ratios of incidence rates per 100k of different age groups in Germany.
# created using matplotlib 3.4.2, pandas 1.3.2, openpyxl 3.0.7
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import datetime
from babel.dates import format_date
import pandas as pd
import sys
import os
import xml.etree.ElementTree as ET
from io import BytesIO

# script options
lang = 'en' if '--lang=en' in sys.argv else 'de_DE'
detailed = '--all' in sys.argv  # include all age groups

# file names
PATH = '.'
INPUT = 'Altersverteilung.xlsx'
OUTPUT = 'altersverteilung%s' % ('_detailed' if detailed else '')  # .png and .svg

# xlsx file must be downloaded from
# https://www.rki.de/DE/Content/InfAZ/N/Neuartiges_Coronavirus/Daten/Altersverteilung.xlsx?__blob=publicationFile
df = pd.read_excel(os.path.join(PATH, INPUT), sheet_name=1, index_col=0).swapaxes(0, 1)  # sheet name 0 or 1

########################################################################
average = df['Gesamt']
sortedIndices = list(reversed(df.columns))[:-1]  # drops Gesamt
assert sortedIndices[0] == '0 - 4' and sortedIndices[-1] == '90+'

assert df.index[0] == '2020_10'
day0 = datetime.datetime(2020, 3, 8, 12, 0, 0)  # last day of first calendar week
days = [day0 + datetime.timedelta(weeks=int(j)) for j in range(df.shape[0])]

topoheights = [0, 5, 10, 25, 50, 100, 200, 400]
topocolors = [mpl.colors.hsv_to_rgb((1, 0, 0.95 - 0.05 * j)) for j in range(len(topoheights))]
topodata = [h / average for h in topoheights]
colors = mpl.cm.tab20b.colors
startweek = 0

fig, ax = plt.subplots(figsize=(12.8, 7.2))
ax.set_xlim(days[startweek], days[-1])
ymax = 3.05
ax.set_ylim(0, ymax)
ax.set_facecolor('black')  # background color, so we do not accidentally leave empty space outside shaded regions
polys = []
for j in range(len(topodata)):
    polys.append(ax.fill_between(days, topodata[0] + ymax, topodata[j],
                    where=topodata[j]<=ymax,  # clipping the polygons at the top gives better SVGs
                    interpolate=True,  # for correct intersection points with ymax
                    linewidth=0, edgecolor=topocolors[j-1] if j%3==2 else None,
                    hatch='\\'*5 if j%6==2 else '/////\\\\' if j%3==2 else None,
                    zorder=-len(topodata)+j, label='≥ %s' % topoheights[j],
                    facecolor=topocolors[j]))
curves = []
for j, idx in reversed(list(enumerate(sortedIndices))[::1 if detailed else 2]):
    curves.extend(ax.plot(days, df[idx] / average, color=colors[j],
                          gid='gid_line',
                          label=idx.replace(' - ', '–'), linewidth=2.2,
                          linestyle=['-','--',':'][j%3]))
curves.extend(ax.plot([days[0], days[-1]], [1,1], label='overall' if lang=='en' else 'gesamt',
                      linewidth=1.5, color='black', zorder=0.9))
plt.grid(color='white', alpha=0.6)
legend2 = ax.legend(curves, [c.get_label() for c in curves], bbox_to_anchor=(1,-0.01), loc="lower left",
                    fontsize=8.2 if detailed else None,
                    title='Age groups' if lang=='en' else 'Altersgruppen', frameon=False)
legend1 = ax.legend(reversed(polys), reversed([p.get_label() for p in polys]), bbox_to_anchor=(1,1.01), loc="upper left",
                    fontsize=8.2 if detailed else None,
                    title='7-day incidence\n(per 100,000)' if lang=='en' else '7-Tage-Inzidenz\n(je 100.000)', frameon=False)
plt.gca().add_artist(legend2)  # as legend1 is wider, we add it last so that it does not get cut off
ax.xaxis.set_minor_locator(mdates.MonthLocator())
for label in ax.xaxis.get_ticklabels():
    label.set_horizontalalignment('left')
ax.text(0, 1.10, 'COVID-19 in Germany by age' if lang=='en' else 'COVID-19-Inzidenz in Deutschland nach Lebensalter',
        transform=ax.transAxes, size=16, weight=600, ha='left')
fmtdate = lambda d: format_date(d, locale=lang)
ax.text(0, 1.021, 'Rate ratio of weekly incidence rate among selected age groups to weekly incidence rate of overall population'
            f'\nin the period from {fmtdate(days[startweek])} to {fmtdate(days[-1])}' if lang=='en' else
            'Verhältnis zwischen 7-Tage-Inzidenz verschiedener Altersgruppen und 7-Tage-Inzidenz der Gesamtbevölkerung'
            f'\nim Zeitraum {fmtdate(days[startweek])}–{fmtdate(days[-1])}', transform=ax.transAxes, size=12, color='#777777',
            linespacing=1.5)
ax.set_ylabel('Rate ratio' if lang=='en' else 'Verhältnis zur 7-Tage-Inzidenz der Gesamtbevölkerung', fontstyle='italic')
ax.set_xlabel('Date' if lang=='en' else 'Datum', fontstyle='italic')
ax.set_axisbelow(True)
ax.text(1, 1.007, 'CC BY-SA · Data source: Altersverteilung Robert Koch-Institut (RKI), dl-de/by-2-0, own represent.'
        if lang=='en' else 'CC BY-SA · Datenquelle: Altersverteilung Robert Koch-Institut (RKI), dl-de/by-2-0, eig. Darst.',
        transform=ax.transAxes, ha='right', fontstretch='condensed', size=9, color='#777777')
plt.tight_layout()

fig.savefig(os.path.join(PATH, f"{OUTPUT}.png"), transparent=False)
fig.savefig(os.path.join(PATH, f"{OUTPUT}.svg"), transparent=True)