import numpy as np
import matplotlib.pyplot as plt
from sklearn import neighbors
import gradio as gr

def train_and_plot(weights, n_neighbors):
    np.random.seed(0)
    X = np.sort(5 * np.random.rand(40, 1), axis=0)
    T = np.linspace(0, 5, 500)[:, np.newaxis]
    y = np.sin(X).ravel()

    # Add noise to targets
    y[::5] += 1 * (0.5 - np.random.rand(8))

    knn = neighbors.KNeighborsRegressor(n_neighbors, weights=weights)
    fit = knn.fit(X, y)
    y_ = knn.predict(T)
    score = knn.score(T, y_)

    plt.figure()
    plt.scatter(X, y, color="darkorange", label="data")
    plt.plot(T, y_, color="navy", label="prediction")
    plt.axis("tight")
    plt.legend()
    plt.title("KNeighborsRegressor (k = %i, weights = '%s')" % (n_neighbors, weights))    

    plt.tight_layout()
    return plt, score


with gr.Blocks() as demo:
    link = "https://scikit-learn.org/stable/auto_examples/neighbors/plot_regression.html#sphx-glr-auto-examples-neighbors-plot-regression-py"
    gr.Markdown("## Nearest Neighbors Regression")
    gr.Markdown(f"This demo is based on this [scikit-learn example]({link}).")
    gr.HTML("<hr>")
    gr.Markdown("In this demo, we learn a noise-infused sine function using k-Nearest Neighbor and observe how the function learned varies as we change the following hyperparameters:")
    gr.Markdown("""1. Weight function
    2. Number of neighbors""")            
    
    with gr.Row():
        weights = gr.Radio(['uniform', "distance"], label="Weights", info="Choose the weight function")
        n_neighbors = gr.Slider(label="Neighbors", info="Choose the number of neighbors", minimum =1, maximum=15, step=1) 
                
    
    with gr.Row():
        with gr.Column(scale=2):
            plot = gr.Plot(label="KNeighborsRegressor Plot")
        with gr.Column(scale=1):
            num = gr.Textbox(label="Test Accuracy")

        
    weights.change(train_and_plot,  inputs=[weights, n_neighbors], outputs=[plot, num])
    n_neighbors.change(train_and_plot,  inputs=[weights, n_neighbors], outputs=[plot, num])
    

if __name__ == "__main__":
    demo.launch()