# Vaš kôd ovdje
X = make_instances(-5,5,1000)
x_train500, x_test500 = model_selection.train_test_split(X, test_size=0.5)

x_train167, x_train333 = model_selection.train_test_split(x_train500, test_size=2/3)
x_test167, x_test333 = model_selection.train_test_split(x_test500, test_size=2/3)


#############################################

y_train167_1 = list(make_labels(x_train167, f, 100))
y_train167_2 = list(make_labels(x_train167, f, 200))
y_train167_5 = list(make_labels(x_train167, f, 500))

y_train333_1 = list(make_labels(x_train333, f, 100))
y_train333_2 = list(make_labels(x_train333, f, 200))
y_train333_5 = list(make_labels(x_train333, f, 500))

y_train500_1 = list(make_labels(x_train500, f, 100))
y_train500_2 = list(make_labels(x_train500, f, 200))
y_train500_5 = list(make_labels(x_train500, f, 500))

##############################################

y_test167_1 = list(make_labels(x_test167, f, 100))
y_test167_2 = list(make_labels(x_test167, f, 200))
y_test167_5 = list(make_labels(x_test167, f, 500))

y_test333_1 = list(make_labels(x_test333, f, 100))
y_test333_2 = list(make_labels(x_test333, f, 200))
y_test333_5 = list(make_labels(x_test333, f, 500))

y_test500_1 = list(make_labels(x_test500, f, 100))
y_test500_2 = list(make_labels(x_test500, f, 200))
y_test500_5 = list(make_labels(x_test500, f, 500))

##############################################

x_train_all = [x_train167, x_train333, x_train500]
y_train_all = [y_train167_1,y_train167_2,y_train167_5,
               y_train333_1,y_train333_2,y_train333_5,
               y_train500_1,y_train500_2,y_train500_5]

x_test_all = [x_test167, x_test333, x_test500]
y_test_all = [y_test167_1,y_test167_2,y_test167_5,
              y_test333_1,y_test333_2,y_test333_5,
              y_test500_1,y_test500_2,y_test500_5]

##############################################

D = range(1,21)

_, plots = plt.subplots(3,3)

for i in range(0,9):
    E_train=[]
    E_test=[]
    i_=[]
    if i<3:
        j=0
    elif i<6:
        j=1
    else:
        j=2
    for d in D:
        i_.append(i)
        poly=preprocessing.PolynomialFeatures(d)
        fi_train = poly.fit_transform(x_train_all[j])
        fi_test = poly.fit_transform(x_test_all[j])

        regr_train = LinearRegression().fit(fi_train,y_train_all[i])

        y_pred_train = regr_train.predict(fi_train)
        y_pred_test = regr_train.predict(fi_test)

        E = mean_squared_error(y_train_all[i], y_pred_train)
        E_train.append(E)
        E = mean_squared_error(y_test_all[i], y_pred_test)
        E_test.append(E)

    plots[i//3][i%3].plot(D, np.log(E_train), label = 'train')
    plots[i//3][i%3].plot(D, np.log(E_test), label = 'test')

    plots[i//3][i%3].legend(loc='best')