"""
Plot total cases timeshiftet to first death
"""
import matplotlib.pyplot as pp
import numpy as np
name="normalized_to_100_cases"

def plot(data, countries):
    figsize = (10,5)
    for loc in data:
        if loc not in countries:
            continue
        time, new_cases, new_deaths, total_cases, total_deaths = data[loc]
        
        pp.figure(name, figsize=figsize)
        day_of_100_cases = np.argwhere(np.array(total_cases) > 99)[0][0]
        new_time_axis = np.arange(len(time)) - day_of_100_cases
        pp.plot(new_time_axis, np.array(total_cases), label=f"{loc}", marker=".")

    pp.yscale("log")
    pp.xticks(rotation=45)
    pp.ylabel("total cases")
    pp.xlabel("days since 100 reported cases")
    pp.legend(frameon=False)
    axis = pp.axis()
    pp.axis([-2, axis[1], 80, axis[3]])

    pos = [(10, 5e5), (30, 5e5), (50, 5e5), (70, 1e3)]
    for i, n in enumerate([1,2,3,7]):
        pp.plot([0, axis[1]], [100, np.exp(np.log(2) / n * axis[1])], color="grey", linestyle="--")
        pp.text(pos[i][0], pos[i][1], f"doubling every\n{n} days", fontsize=8)

    pp.tight_layout()

    pp.savefig(name+".png")