ax1.set_ylabel("traffic")

ax2 = ax1.twinx()
ax2.set_ylabel("cases")
ax2.plot(days, traffic)
plt.show()
##########################

# Choose optimal KFold
print("Performing Cross Validation...")

cross_validate.k_folds_cross_validation(
    1, cases, traffic, days, pred_type='cases', model_type='knn', Q=1, K=18, C='N/A')

# Choose optimal polynomial features
cross_validate.poly_feature_cross_validation(
    2, cases, traffic, days, pred_type='cases', model_type='knn', folds=2, K=18, C='N/A')

# Choose optimal neighbours for KNN = 18
cross_validate.knn_cross_validation(
    3, cases, traffic, days, pred_type='cases', model_type='knn', folds=2, Q=1, C='N/A')

# TRAFFIC ==> CASES
kf = KFold(n_splits=2)
pred_array = []
y = []
plt.figure(5)
plt.plot(days, cases)
for train, test in kf.split(traffic):
    model = KNeighborsRegressor(n_neighbors=18).fit(
        traffic[train], cases[train])
    predictions = model.predict(traffic[test])
cases = data.iloc[:, 1]
traffic = data.iloc[:, 2]

cases_df = cases
traffic_df = traffic
cases = pd.DataFrame(cases).to_numpy()
traffic = pd.DataFrame(traffic).to_numpy()
cases = cases.reshape(-1, 1)

# Choose optimal KFold
print("Performing Cross Validation...")
cross_validate.k_folds_cross_validation(
    1, cases, traffic, days, pred_type='cases', model_type='lasso', Q=1, K='N/A', C=10)

# Choose optimal polynomial features
cross_validate.poly_feature_cross_validation(
    2, cases, traffic, days, pred_type='cases', model_type='lasso', folds=2, K='N/A', C=10)

# Choose optimal value for C penalty
cross_validate.c_penalty_cross_validation(
    3, cases, traffic, days, pred_type='cases', model_type='lasso', folds=2, Q=5, K='N/A')

# TRAFFIC ==> CASES
kf = KFold(n_splits=5)
y = []
p = []
plt.figure(5)
plt.plot(days, cases)
for train, test in kf.split(traffic):
    a = 1/2*10
    model = Lasso(alpha=a).fit(traffic[train], cases[train])
    predictions = model.predict(traffic[test])