import numpy as np

import matplotlib.pyplot as plt

from sklearn import datasets

from sklearn.model_selection import train_test_split

from sklearn.linear_model import LinearRegression

from sklearn.metrics import mean_squared_error, r2_score

diabetes = datasets.load_diabetes()

# Para ver a descrição do dataset

print(diabetes.DESCR)

# Visualizar os dados

X = diabetes.data

y = diabetes.target

# Como temos múltiplas características, vamos usar apenas a primeira para simplificar o exemplo

X = X[:, np.newaxis, 2]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = LinearRegression()

model.fit(X_train, y_train)

y_pred = model.predict(X_test)

# Avaliar o desempenho do modelo

print('Coefficients: \n', model.coef_)

print('Mean squared error: %.2f' % mean_squared_error(y_test, y_pred))

print('Coefficient of determination: %.2f' % r2_score(y_test, y_pred))