"""
Plot delay of countries
"""
import matplotlib.pyplot as pp
import numpy as np
name="delay"

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_above_hundred_cases = np.argwhere(np.array(total_cases) > 100)[0][0]
        
        new_time_axis = np.arange(len(time)) - day_of_above_hundred_cases
        pp.plot(new_time_axis, np.array(total_cases), label=f"{loc} - {day_of_above_hundred_cases}", marker=".")

    pp.yscale("log")
    pp.xticks(rotation=45)
    pp.legend(frameon=False)
    pp.tight_layout()

    pp.savefig(name+".png")

    # plot delay
    pp.clf()
    delay = {}
    tcases = [50, 100, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 80500,81000,81500,82000,83000,835000,84000,84500] + list(np.arange(10000, np.max(data["United States"][3]), 2500))
    tcases.sort()

    time_us = data["United States"][0]
    try:    
        case_date_us = np.array([time_us[np.argwhere(np.array(data["United States"][3]) > cases)[0][0]] for cases in tcases])
    except:
        print("workaround delay_us:40")
        case_date_us = []
        usdat = np.array(data["United States"][3])
        for cases in tcases:
            try:
                casepos = np.argwhere(usdat > cases)[0][0]
                casetime = time_us[casepos]
            except:
                print(f"case date for {cases} not found")
                casetime = np.nan
                #continue
            case_date_us.append(casetime)
        case_date_us = np.array(case_date_us)
            
    total_cases_us = data["United States"][3]
    for loc in data:
        if loc not in countries:
            continue
        if loc in ["United States", "World"]:
            continue
        this_delay = []
        for i, cases in enumerate(tcases):
            try:
                case_date = data[loc][0][np.argwhere(np.array(data[loc][3]) > cases)[0][0]]
            except:
                this_delay.append(np.nan)
                continue
            #case_date = np.array([time_us[np.argwhere(np.array(data[loc][3]) > cases)[0][0]] for cases in tcases])
        
            this_delay.append((case_date_us[i] - case_date).days)

        delay[loc] = this_delay

    sort = np.argsort(delay)
    
    for loc in delay:
        pp.plot(tcases, delay[loc], label=loc, marker=".")
    #pp.plot(np.arange(len(delay)), np.array(delay)[sort] - np.min(delay), marker="o", linestyle="")
    #pp.xticks(ticks=np.arange(len(delay)), labels=np.array(d_loc)[sort])
    pp.xticks(rotation=45)
    pp.legend(frameon=False)
    pp.xlabel("total cases")
    pp.ylabel("delay to United States [days]")
    pp.tight_layout()
    pp.savefig(name+"_usa.png")