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

def plot(data, countries):

    for loc in data:
        if loc not in countries:
            continue
        time, new_cases, new_deaths, total_cases, total_deaths = data[loc]
        
        # total cases
        pp.figure("total_cases")
        pp.plot(time, total_cases, label=f"{loc}")

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

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

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

    for name in ["total_cases", "total_deaths", "new_cases", "new_deaths"]:
        postprocess(name)

def postprocess(name):
    pp.yscale("log")
    pp.xticks(rotation=90)
    pp.legend(frameon=False)
    pp.tight_layout()

    pp.savefig(f"{name}.png")