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

def plot(data, countries):
    figsize = (10,5)

    # plot delay
    pp.clf()
    delay = {}
    tcases = [50, 100, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000, 12500, 15000, 17500,
            20000, 22500, 25000, 27500,
            30000, 32500, 35000, 37500,
            40000, 42500, 45000, 47500,
            50000, 52500, 55000, 57500,
            60000, 62500, 65000, 67500,
            70000, 72500, 75000, 77500]+list(np.arange(80000,max(data["China"][3]),500))

    time_china = data["China"][0]
    case_date_china = np.array([time_china[np.argwhere(np.array(data["China"][3]) > cases)[0][0]] for cases in tcases])
    total_cases_china = data["China"][3]
    for loc in data:
        if loc not in countries:
            continue
        if loc in ["China", "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_china[np.argwhere(np.array(data[loc][3]) > cases)[0][0]] for cases in tcases])
        
            this_delay.append((case_date_china[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 china [days]")
    pp.tight_layout()
    pp.savefig(name+"_china.png")