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

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

    pp.yscale("log")
    pp.ylabel("total cases")
    pp.xlabel("day since first reported death")
    pp.xticks(rotation=45)
    pp.legend(frameon=False)
    pp.tight_layout()

    pp.savefig(name+".png")