import gradio as gr
import numpy as np
from sklearn.linear_model import LinearRegression
import spaces

# Define the linear regression function
def linear_regression(X, y, new_data):
    X = np.array(X).reshape(-1, 1)  # Reshape for sklearn
    y = np.array(y)

    # Perform the computation within the GPU context
    with spaces.GPU():
        model = LinearRegression()
        model.fit(X, y)
        prediction = model.predict(np.array(new_data).reshape(-1, 1))

    return prediction.tolist()

# Create the Gradio interface
with gr.Blocks(title="Linear Regression with ZeroGPU") as iface:
    gr.Markdown("# Linear Regression Example")
    gr.Markdown("This example performs linear regression using ZeroGPU for computation.")
    
    with gr.Row():
        with gr.Column():
            X = gr.Textbox(lines=5, label="Input Features (X)", placeholder="e.g., 1, 2, 3, 4")
            y = gr.Textbox(lines=5, label="Target Values (y)", placeholder="e.g., 2, 4, 6, 8")
            new_data = gr.Textbox(lines=1, label="New Data for Prediction", placeholder="e.g., 5")
            submit_btn = gr.Button("Predict")
        
        output = gr.Textbox(lines=5, label="Predicted Values")
    
    submit_btn.click(
        fn=linear_regression,
        inputs=[X, y, new_data],
        outputs=output
    )

# Launch the interface
iface.launch(server_name="0.0.0.0", server_port=7860)