import gradio as gr

import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets, linear_model
from sklearn.metrics import mean_squared_error, r2_score

FIGSIZE = (10,10)

feature_names = ["age", "body-mass index (BMI)", "blood pressure", 
                "total serum cholesterol", "low-density lipoproteins (LDL)", 
                "high-density lipoproteins (HDL)", "total cholesterol / HDL", 
                "log of serum triglycerides level (possibly)","blood sugar level"]

def create_dataset(feature_id=2):
    # Load the diabetes dataset
    diabetes_X, diabetes_y = datasets.load_diabetes(return_X_y=True)

    # Use only one feature
    diabetes_X = diabetes_X[:, np.newaxis, feature_id]

    # Split the data into training/testing sets
    diabetes_X_train = diabetes_X[:-20]
    diabetes_X_test = diabetes_X[-20:]

    # Split the targets into training/testing sets
    diabetes_y_train = diabetes_y[:-20]
    diabetes_y_test = diabetes_y[-20:]

    return diabetes_X_train, diabetes_X_test, diabetes_y_train, diabetes_y_test

def train_model(input_data):

    # We removed the sex variable 
    if input_data == 'age':
        feature_id = 0
    else:
        feature_id = feature_names.index(input_data) + 1

    diabetes_X_train, diabetes_X_test, diabetes_y_train, diabetes_y_test = create_dataset(feature_id)

    
    # Create linear regression object
    regr = linear_model.LinearRegression()

    # Train the model using the training sets
    regr.fit(diabetes_X_train, diabetes_y_train)

    # Make predictions using the testing set
    diabetes_y_pred = regr.predict(diabetes_X_test)

    mse = mean_squared_error(diabetes_y_test, diabetes_y_pred)
    r2 = r2_score(diabetes_y_test, diabetes_y_pred)

    # Plot outputs
    fig = plt.figure(figsize=FIGSIZE)
    
    # plt.title(input_data)
    plt.scatter(diabetes_X_test, diabetes_y_test, color="black")
    plt.plot(diabetes_X_test, diabetes_y_pred, color="blue", linewidth=3)

    plt.xticks(())
    plt.yticks(())



    return fig, regr.coef_, mse, r2

title = "Linear Regression Example 📈"
description = "The example shows how linear regression attempts to draw a straight line that will best minimize the residual sum of squares between the observed responses in the dataset"
with gr.Blocks() as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description)

    with gr.Column():
    
        with gr.Row():
            plot = gr.Plot()
            with gr.Column():
                input_data = gr.Dropdown(choices=feature_names, label="Feature", value="body-mass index")
                coef = gr.Textbox(label="Coefficients")
                mse = gr.Textbox(label="MSE")
                r2 = gr.Textbox(label="R2")

        input_data.change(fn=train_model, inputs=[input_data], outputs=[plot, coef, mse, r2], queue=False)

    
demo.launch(enable_queue=True)