"""
Plot total cases of countries over time on log scale
"""
import matplotlib.pyplot as pp
import numpy as np

def plot(data, countries, pop, **kwargs):
    figsize = (10,5)
    tcp, tc = pp.subplots(figsize=figsize)  # total cases
    tdp, td = pp.subplots(figsize=figsize)  # total deaths
    tip, ti = pp.subplots(figsize=figsize)  # total (currently) infected
    ncp, nc = pp.subplots(figsize=figsize)  # new cases
    ndp, nd = pp.subplots(figsize=figsize)  # new deaths
    for loc in data:
        if loc not in countries:
            continue
        time, new_cases, new_deaths, total_cases, total_deaths = data[loc]
        
        # total cases
        tc.plot(time, total_cases, label=f"{loc}", marker=".")

        # total deaths
        td.plot(time, total_deaths, label=f"{loc}", marker=".")

        # new cases
        nc.plot(time, new_cases, label=f"{loc}", marker=".")

        # new deaths
        nd.plot(time, new_deaths, label=f"{loc}", marker=".")

        # currently infected
        delay = 21
        current_infected = np.array(total_cases[delay:]) - np.array(total_deaths[:-delay]) - np.array(total_cases[:-delay])

        ti.plot(time[:-delay], current_infected, label=f"{loc}", marker=".")

    for ax, fig, name in [(tc, tcp, "total_cases"), (td, tdp, "total_deaths"), (nc, ncp, "new_cases"), (nd, ndp, "new_deaths"), (ti, tip, "current_infected")]:
        ax.set_yscale("log")
        ax.set_ylabel(name)
        for tick in ax.get_xticklabels():
            tick.set_rotation(45)
        ax.legend(frameon=False)
        ax.grid(True)
        fig.tight_layout()
    

        fig.savefig(name+".png")