import gradio as gr
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputRegressor
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

def compare(max_depth,n_estimators):
    rng = np.random.RandomState(1)
    X = np.sort(200 * rng.rand(600, 1) - 100, axis=0)
    y = np.array([np.pi * np.sin(X).ravel(), np.pi * np.cos(X).ravel()]).T
    y += 0.5 - rng.rand(*y.shape)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, train_size=400, test_size=200, random_state=4
    )

    regr_multirf = MultiOutputRegressor(
        RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth, random_state=0)
    )
    regr_multirf.fit(X_train, y_train)

    regr_rf = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth, random_state=2)
    regr_rf.fit(X_train, y_train)

    # Predict on new data
    y_multirf = regr_multirf.predict(X_test)
    y_rf = regr_rf.predict(X_test)

    # Plot the results
    fig, ax = plt.subplots()
    s = 50
    a = 0.4
    ax.scatter(
        y_test[:, 0],
        y_test[:, 1],
        edgecolor="k",
        c="navy",
        s=s,
        marker="s",
        alpha=a,
        label="Data",
    )
    ax.scatter(
        y_multirf[:, 0],
        y_multirf[:, 1],
        edgecolor="k",
        c="cornflowerblue",
        s=s,
        alpha=a,
        label="Multi RF score=%.2f" % regr_multirf.score(X_test, y_test),
    )
    ax.scatter(
        y_rf[:, 0],
        y_rf[:, 1],
        edgecolor="k",
        c="c",
        s=s,
        marker="^",
        alpha=a,
        label="RF score=%.2f" % regr_rf.score(X_test, y_test),
    )
    ax.set_xlim([-6, 6])
    ax.set_ylim([-6, 6])
    ax.set_xlabel("target 1")
    ax.set_ylabel("target 2")
    ax.set_title("Comparing random forests and the multi-output meta estimator")
    ax.legend()
    return fig

title = "Comparing random forests and the multi-output meta estimator"
with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown("This app demonstrates random forests and the multi-output meta estimator comparison")

    max_depth = gr.Slider(minimum=10, maximum=50, step=1, label = "Maximum Depth")
    n_estimators = gr.Slider(minimum=50, maximum=300, step=1, label = "Number of Estimators")

    plot = gr.Plot(label=title)
    n_estimators.change(fn=compare, inputs=[max_depth,n_estimators], outputs=[plot])
    max_depth.change(fn=compare, inputs=[max_depth,n_estimators], outputs=[plot])

demo.launch()