import gradio as gr
import pandas as pd

from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.datasets import load_diabetes
from sklearn.multioutput import MultiOutputRegressor

# Assume y is now a DataFrame with multiple columns
X, y = load_diabetes(return_X_y=True, as_frame=True)
y = pd.DataFrame([y, y]).T

# Make the estimator multi-output capable
est = MultiOutputRegressor(HistGradientBoostingRegressor()).fit(X, y)


def predict(input_df):
    prediction = est.predict(input_df)
    # Assume y has columns named "target1", "target2", etc.
    return pd.DataFrame(prediction, columns=y.columns)


iface = gr.Interface(
    fn=predict,
    inputs=gr.Dataframe(
        value=X.head(1),
        headers=list(X.columns),
        col_count=(X.shape[1], "fixed"),
        row_count=(1, "dynamic"),
        datatype=X.dtypes.apply(str).replace("float64", "number").values.tolist(),
    ),
    outputs=gr.Dataframe(
        value=y.head(1),
        headers=list(y.columns),
        col_count=(y.shape[1], "fixed"),
        datatype=y.dtypes.apply(str).replace("float64", "number").values.tolist(),
    ),
)
iface.launch()