"""
Plot overview plot for each country separately
"""
import matplotlib.pyplot as pp
from mpl_toolkits.axisartist.parasite_axes import HostAxes, ParasiteAxes
import numpy as np
import time as time_module
import pickle
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
basename="all_"

# manual y adjustments for new cases
corr = {"Chile": 10000,
        "China": 4000,
        "Ecuador": 3000,
        "Kazakhstan": 5000,
        "Kyrgyzstan": 2200,
        "Peru": 15000,
        "Bolivia": 3000,
        "Yemen": 220,
        "Turkey": 100000,
        "Spain": 40000,
}

def plot(data, countries, pop, metadata={}, **kwargs):
    figsize = (10,5)
    
    vaccs = []

    for loc in data:
        try:
            if loc == "International":
                continue
            name = basename+loc
            time, new_cases, new_deaths, total_cases, total_deaths, total_vaccinations, stringency_index, new_vaccinations = data[loc]['time'], data[loc]['new_cases'], data[loc]['new_deaths'], data[loc]['total_cases'], data[loc]['total_deaths'], data[loc]['total_vaccinations'], data[loc]['stringency_index'], data[loc]['new_vaccinations']
            people_fully_vaccinated = data[loc]['people_fully_vaccinated']

            fig, ax1 = pp.subplots(num=name, figsize=figsize)


            ax2 = ax1.twinx()
            ax2.plot(time, np.array(total_deaths)*10, label="Total deaths (x10)", marker="", linestyle="--", color="green")
            ax2.plot(time, total_cases, label=f"Total cases", marker="", linestyle="-", color="blue")

            ax1.plot(time, np.array(new_deaths)*10, label="raw new deaths (x10)", color="grey", linestyle=":")
            ax1.plot(time[3:-3], np.convolve(new_deaths, np.ones((7,))/7, mode="valid")*10, label="new deaths 7day mean (x10)", color="black", linestyle="--", linewidth=2)

            ax1.plot(time, new_cases, label="raw new cases", color="grey", linestyle="-")
            ax1.plot(time[3:-3], np.convolve(new_cases, np.ones((7,))/7, mode="valid"), label="new cases 7day mean", color="orange", linestyle="-", linewidth=2)

            # plot vaccinations
            if not np.isnan(total_vaccinations).all():
                # notify of new vaccine programs
                if np.isnan(total_vaccinations[-2]):
                    print(f"{loc} starts vaccinating, adding to plot")

                ax2.plot(time, np.array(total_vaccinations), label=f"Total vaccination doses", marker="", linestyle="-.", color="crimson")
                ax2.plot(time, np.array(people_fully_vaccinated), label="fully vaccinated", marker="", linestyle="-", color="crimson")

            # fix lower bound of plot
            for ax in (ax1, ax2):
                axis = ax.axis()
                if ax is ax1: # adjust left (new cases) axes upper boundary for given countries
                    if loc in corr:
                        axis = [axis[0], axis[1], axis[2], corr[loc]]
                ax.axis([axis[0], axis[1], -1, axis[3]])

            # fix population
            #try:
            #    print(loc, pop[loc]['pop'] - metadata[loc]['population'])
            #except:
            #    pop[loc]['pop'] = metadata[loc]['population']

            # if we know population: plot 500 new cases / 1million inhabitants as a rough measure for comparison
            # also set color for infection level indicator
            infection_level_indicator = "grey"
            try:
            #if False:
                warn_thresh = 500e-6 * pop[loc]['pop']/7
                info_thresh = 50e-6 * pop[loc]['pop']/7
                low_thresh = 5e-6 * pop[loc]['pop']/7
                actual_level = np.mean(new_cases[-7:])
                infection_level_indicator = "green"
                if actual_level > low_thresh:
                    infection_level_indicator = "gold"
                if actual_level > info_thresh:
                    infection_level_indicator = "peru"
                if actual_level > warn_thresh:
                    infection_level_indicator = "r"
                bounds = ax1.axis()
                # for small population numbers, give floats instead of ints
                if pop[loc]['pop'] < 5e6:
                    ax1.plot([bounds[0], bounds[1]], [warn_thresh]*2, color="red", linestyle=":", label=f"500 new cases / week / 1M inh.: {warn_thresh:1.2f}".replace(",", "."))
                    ax1.plot([bounds[0], bounds[1]], [info_thresh]*2, color="peru", linestyle=":", label=f"50 new cases / week / 1M inh.: {info_thresh:1.2f}".replace(",", "."))
                    ax1.plot([bounds[0], bounds[1]], [low_thresh]*2, color="gold", linestyle=":", label=f"5 new cases / week / 1M inh.: {low_thresh:1.2f}".replace(",", "."))
                else: # but not for big populations
                    ax1.plot([bounds[0], bounds[1]], [warn_thresh]*2, color="red", linestyle=":", label=f"500 new cases / week / 1M inh.: {int(warn_thresh):,}".replace(",", "."))
                    ax1.plot([bounds[0], bounds[1]], [info_thresh]*2, color="peru", linestyle=":", label=f"50 new cases / week / 1M inh.: {int(info_thresh):,}".replace(",", "."))
                    ax1.plot([bounds[0], bounds[1]], [low_thresh]*2, color="gold", linestyle=":", label=f"5 new cases / week / 1M inh.: {int(low_thresh):,}".replace(",", "."))

                ax1.axis(bounds)
            except:
                print(f"=====> population unknown for {loc}, skipping plot enhancements")
                

            # stringency of countermeasures as background contourf
            axbounds = ax1.axis()
            ax1.contourf(time, [-1, 1e10], [stringency_index]*2, cmap="Greys", alpha=0.3, levels=99)
            ax1.axis(axbounds)


            # disabled for now: put second total-cases-per-million-inhabitants axis besides total-cases-axis
            if loc in pop and False:
                # according to https://matplotlib.org/3.2.2/gallery/axisartist/demo_parasite_axes.html
                host1 = HostAxes(ax1, [0.15, 0.1, 0.65, 0.8])
                par1 = ParasiteAxes(host1, sharex=host1)
                host1.append(par1)
                host2 = HostAxes(ax2, [0.15, 0.1, 0.65, 0.8])
                par2 = ParasiteAxes(host2, sharex=host1)
                host2.append(par2)
                ax3 = ax1.twinx()
                ax3.plot(time, total_cases*1e6/pop[loc]['pop'], linestyle="", marker="")
                ax3.set_ylabel("total cases per 1M inhabitants")

            # plot infection level indicator
            ax1.annotate('Infection status:', xy=(0.85, 1.02), xycoords="axes fraction")
            ax1.annotate('•', xy=(0.99, 1.02), xycoords="axes fraction", color=infection_level_indicator, fontsize="x-large")
            #ax1.xticks(rotation=45)
            #ax1.set_xlabel("date")
            ax1.set_ylabel("new cases")
            ax2.set_ylabel("total cases")
            fig.legend(frameon=False, loc="upper left", bbox_to_anchor=(0,1), bbox_transform=ax1.transAxes)
            title = loc
            if loc in pop:
                #pp.title(f"{loc}", population = "+f"{pop[loc]['pop']:,}".replace(",","."))
                title += ", population = "+f"{pop[loc]['pop']:,}".replace(",",".")
            if not np.isnan(people_fully_vaccinated[-1]):
                title += ", vac rate: "+f"{people_fully_vaccinated[-1]/pop[loc]['pop']*100:1.3f}%"
                vaccs.append([loc, people_fully_vaccinated[-1], people_fully_vaccinated[-1]/pop[loc]['pop']*100]) # bookkeeping for overview
            ax1.set_title(title)
            fig.tight_layout()
            pp.text(0.002,0.005, f"plot generated {time_module.strftime('%Y-%m-%d %H:%M')}, CC-by-sa-nc, origin: dukun.de/corona, datasource: ourworldindata.org/coronavirus-source-data", color="dimgrey", fontsize=8, transform=fig.transFigure)

            pp.savefig("img/ac_"+name+".png")
            pp.close(fig)
        except Exception as e:
            print(f"=====> plotting failed for {loc}, skipping plot. Error: {e}")
    
    ## vaccination overview html
    with open("index.html", "w") as f:
        # site header
        with open("index.html.head", "r") as g:
            f.write(g.read())
        # table header
        f.write("<table><tr><th>Land</th><th>Impfungen</th><th>Impfrate</th><th>Impfstoffe</th></tr>\n")
        # data
        for loc, tvac, rvac in vaccs:
            line =  f"<tr><td>{loc}</td><td>" + \
                    f"{int(tvac):,d}".replace(",",".") + \
                    f"</td><td>{rvac:3.3f}%</td>".replace(".", ",")
            if "vaccines" in metadata[loc]:
                line += f"<td>{metadata[loc]['vaccines']}</td></tr>\n"
            else:
                line += f"<td>-</td></tr>\n"
            f.write(line) 
        # table footer
        f.write("</table>\n")
        # site footer
        with open("index.html.foot", "r") as g:
            f.write(g.read())