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 comparison of random forest with multi-output meta estimator for multi-output regression.
    A random forest regressor is trained on randomly generated data which is used as baseline and compared with multi-output meta estimator trained on the same dataset. 
    The predicted outputs from each model are visualized in the plot together with the actual data. 
    The maximum depth and number of estimator of the random forest can be adjusted and the effect can be seen in the resulting plot.
    This app is developed based on [scikit-learn example](https://scikit-learn.org/stable/auto_examples/ensemble/plot_random_forest_regression_multioutput.html#sphx-glr-auto-examples-ensemble-plot-random-forest-regression-multioutput-py)
    """)

    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()