"""
Experiment summary
------------------
Treat each province/state in a country cases over time
as a vector, do a simple K-Nearest Neighbor between
countries. What country has the most similar trajectory
to a given country?

Plots similar countries
"""

import sys
sys.path.insert(0, '..')

from utils import data
import os
import sklearn
import numpy as np
import json
import matplotlib.pyplot as plt

plt.style.use('fivethirtyeight')

# ------------ HYPERPARAMETERS -------------
BASE_PATH = '../COVID-19/csse_covid_19_data/'
# ------------------------------------------

confirmed = os.path.join(
    BASE_PATH,
    'csse_covid_19_time_series',
    'time_series_covid19_confirmed_global.csv')
confirmed = data.load_csv_data(confirmed)
features = []
targets = []

fig = plt.figure(figsize=(12, 12))
ax = fig.add_subplot(111)
cm = plt.get_cmap('jet')
NUM_COLORS = 0
LINE_STYLES = ['solid', 'dashed', 'dotted']
NUM_STYLES = len(LINE_STYLES)

dist_diff = os.path.join('../exp/results/', 'knn_raw.json')
f = open(dist_diff,)
dist_diff = json.load(f)

for region, dist in dist_diff.items():
    plt.style.use('fivethirtyeight')

    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111)
    cm = plt.get_cmap('jet')

    other_region = dist['manhattan'][0]
    regions = [region, other_region]
    for val in regions:
        df = data.filter_by_attribute(
            confirmed, "Country/Region", val)
        cases, labels = data.get_cases_chronologically(df)
        cases = cases.sum(axis=0)

        lines = ax.plot(cases, label=val)

    ax.set_ylabel('# of confirmed cases')
    ax.set_xlabel("Time (days since Jan 22, 2020)")
    ax.set_yscale('log')
    ax.legend()
    plt.tight_layout()
    region = region.replace('*', '')
    other_region = other_region.replace('*', '')
    plt.title(f'Comparing confirmed cases in {region} and {other_region}')
    plt.savefig(f'results/raw_manhattan/{region}.png')
    plt.close()
    print(region)
[(13, 'sys.path.insert', 'sys.path.insert', False, 'import sys\n'), (22, 'matplotlib.pyplot.style.use', 'plt.style.use', True, 'import matplotlib.pyplot as plt\n'), (28, 'os.path.join', 'os.path.join', False, 'import os\n'), (32, 'utils.data.load_csv_data', 'data.load_csv_data', False, 'from utils import data\n'), (36, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (38, 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', True, 'import matplotlib.pyplot as plt\n'), (43, 'os.path.join', 'os.path.join', False, 'import os\n'), (45, 'json.load', 'json.load', False, 'import json\n'), (48, 'matplotlib.pyplot.style.use', 'plt.style.use', True, 'import matplotlib.pyplot as plt\n'), (50, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (52, 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', True, 'import matplotlib.pyplot as plt\n'), (68, 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', True, 'import matplotlib.pyplot as plt\n'), (71, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (72, 'matplotlib.pyplot.savefig', 'plt.savefig', True, 'import matplotlib.pyplot as plt\n'), (73, 'matplotlib.pyplot.close', 'plt.close', True, 'import matplotlib.pyplot as plt\n'), (57, 'utils.data.filter_by_attribute', 'data.filter_by_attribute', False, 'from utils import data\n'), (59, 'utils.data.get_cases_chronologically', 'data.get_cases_chronologically', False, 'from utils import data\n')]