# Author: Geoffroy Chaussonnet, Karlsruhe Institute of Technology, Germany
# Script to load the data related to the COVID-19 gathered by the Johns Hopkins University and to plot the phase diagram of the confirmed cases or deaths. X-axis is the gradient (=new case/day) and Y-axis is the curvature (=variation of new case/day)
# Source of the data: https://github.com/CSSEGISandData/COVID-19

from pylab import *
from pathlib import Path
from covid_utils import *

############### Basic use #############################
# example: plot_phase_country("US",dataParam,displayParam,fitParam,'3/22/20',ax)
# Argument 1: name of the country, as it appears in the CSV file
# Argument 2: data parameters, AUTOMATICALLY generated
# Argument 3: display parameters, AUTOMATICALLY generated
# Argument 4: fitting parameters, AUTOMATICALLY generated
# Argument 5: date of the confinement start. If no confinement, enter a date in the future
# Argument 6: axis (matplotlib object) where to plot the curves

######################## Definition of Functions (BEGIN) ############################
from covid_utils import file_name


def plot_phase_country(area, data, quar_date, ax, field, smoothing, y_scale):
    print("########## Treating country: %18s ###########" %('{0:^18}'.format(area)))
    filter_date = data['FilterDate']

    # Extract evolution for this country
    curvature = evolution_country(area, data, field,
                                  "smoothedCurvature", filter_date, smoothing)
    gradient = evolution_country(area, data, field, "daily", filter_date,
                                 smoothing)

    # Filter data for more than 100 cases
    cumul = evolution_country(area, data, field, "cumulative",
                              filter_date, smoothing)
    good_data = (cumul > 100)

    gradient = gradient[good_data]
    curvature = curvature[good_data]
    dates = data['Dates'][good_data]
    
    # find the quarantine date 
    quar_date = dateIn(quar_date)
    quar_indices = dates >= quar_date

    # smooth
    scurv = curvature
    sgrad = gradient
    window_length, polyorder = smoothing
    if window_length != 0:
        scurv = savgol_filter(curvature, window_length, polyorder)
        sgrad = savgol_filter(gradient, window_length, polyorder)

    # draw the diagram
    if y_scale == 'log':
        scurv = np.ma.masked_where(scurv <= 0, scurv)
    p = ax.semilogy(sgrad, scurv, ls='-', marker='o', lw=4.0, label=area)
    col = p[0].get_color()
    ax.scatter(sgrad[-1], scurv[-1], c=col, s=100, marker="s")

    if sum(quar_indices) > 0:  # Quarantine found
        # Plot the quarantine date
        ax.scatter(sgrad[quar_indices][0], scurv[quar_indices][0], c=col, s=300,
                   marker="X")


def setDisplayParam(field, zone, figures_path):
    displayParam = {}

    strUnit, txtField = unit_and_field(field)
    txtEvol = "Phase portrait from"

    title_and_y_axis(displayParam, strUnit, txtEvol, txtField,
                     "{} {}\n (Source: Johns Hopkins University)")

    name = file_name(figures_path, "{}_phase_diagram_Covid19_{}_{}_for_{}.png",
                     txtEvol, txtField, zone)
    displayParam['FileName'] = name
    return displayParam


######################## Definition of Functions (END) ############################


############## Execution section ################

def main():
    # Path to the folder containing the time series:
    data_path="https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/"
    figures_path = "../FIGURES"
    start_date = dt.date(2020, 1,1)   # Start date of the plot:

    yscale = 'linear'   # recommended for phase diagram
    #yscale = 'log'

    # Type of value to analyse:
    field = "Confirmed"
    #field = "Deaths"
    #field = "Active"
    #field = "DeathRate"

    # Smoothing: (mandatory for phase diagram)
    smoothing = (9, 3)  # [window size, order of fitting polynomial]

    # Type of zones (see in the execution section)
    zone = "countries"
    #zone = "continents"
    main_plot(data_path, figures_path, field, start_date, smoothing, yscale,
              zone)


def main_plot(data_path, figures_path, field, start_date, smoothing, yscale,
              zone):
    # Initialisation
    ensure_figures_directory_exists(figures_path)
    data = load_data(data_path, start_date=start_date)
    displayParam = setDisplayParam(field, zone, figures_path)
    # Set graphic objects
    close(1)
    fig = figure(num=1, figsize=(10, 6))
    ax = fig.add_subplot(111)
    if zone == "continents":
        areas = ["EU", "China", "US", "Africa"]
    elif zone == "countries":
        areas = ["Italy", "Spain", "Germany", "France", "Korea, South", "US"]
    else:
        areas = ["World"]
    for area in areas:
        quar_date = data['Confinement'].get(area, '1/1/99')
        plot_phase_country(area, data, quar_date, ax,
                           field, smoothing, yscale)
    # Add graph decorations
    ax.set_title(displayParam['title'])
    ax.set_yscale(yscale)
    ax.set_xlabel(r'Gradient [new case/day]')
    ax.set_ylabel(r'Curvature [new case/day$^2$]')
    ax.legend(loc=2)
    ax.grid(which='major', color='grey', linestyle='-', linewidth=1)
    ax.grid(which='minor', color='grey', linestyle='-', linewidth=0.5)
    fig.tight_layout()
    savefig(displayParam['FileName'], dpi=600, bbox='tight')
    show()


if __name__ == "__main__":
    main()
