from sklearn import datasets, ensemble
from sklearn.inspection import permutation_importance
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
import plotly.graph_objs as go
import numpy as np
import plotly.express as px
import pandas as pd


import gradio as gr


diabetes = datasets.load_diabetes(as_frame=True)
X, y = diabetes.data, diabetes.target


def display_table(row_number):
    X, y = diabetes.data, diabetes.target
    XX = pd.concat([X, y], axis=1)
    temp_df = XX[row_number : row_number + 5]
    Statement = f"Displaying rows from row {row_number} to {row_number+5}"
    return Statement, temp_df


def train_model(
    test_split,
    learning_rate,
    n_estimators,
    max_depth,
    min_samples_split,
    loss,
    duration,
):
    X, y = diabetes.data, diabetes.target
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_split, random_state=42
    )
    params = {
        "n_estimators": n_estimators,
        "max_depth": max_depth,
        "min_samples_split": min_samples_split,
        "learning_rate": learning_rate,
        "loss": loss,
    }
    global reg
    reg = ensemble.GradientBoostingRegressor(**params)
    reg.fit(X_train, y_train)

    mse = mean_squared_error(y_test, reg.predict(X_test))

    x = np.arange(params["n_estimators"]) + 1
    train_score = reg.train_score_
    test_score = np.zeros((params["n_estimators"],), dtype=np.float64)
    for i, y_pred in enumerate(reg.staged_predict(X_test)):
        test_score[i] = mean_squared_error(y_test, y_pred)

    test_score = test_score

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=x,
            y=train_score,
            mode="lines",
            name="Training Set Deviance",
            line=dict(color="blue"),
        )
    )
    fig.add_trace(
        go.Scatter(
            x=x,
            y=test_score,
            mode="lines",
            name="Test Set Deviance",
            line=dict(color="red"),
        )
    )

    frames = [
        go.Frame(
            data=[
                go.Scatter(
                    x=x[: k + 1],
                    y=train_score[: k + 1],
                    mode="lines",
                    line=dict(color="blue"),
                ),
                go.Scatter(
                    x=x[: k + 1],
                    y=test_score[: k + 1],
                    mode="lines",
                    line=dict(color="red"),
                ),
            ],
            name=f"frame{k}",
        )
        for k in range(1, len(x))
    ]

    fig.frames = frames

    fig.update_layout(
        title="Deviance",
        xaxis_title="Boosting Iterations",
        yaxis_title="Deviance",
        legend=dict(x=0, y=1),
        updatemenus=[
            dict(
                type="buttons",
                showactive=False,
                direction="right",
                pad={"r": 10},
                buttons=[
                    dict(
                        label="Play",
                        method="animate",
                        args=[
                            None,
                            dict(
                                frame=dict(duration=duration, redraw=True),
                                fromcurrent=True,
                                transition=dict(duration=0),
                            ),
                        ],
                    ),
                    dict(
                        label="Pause",
                        method="animate",
                        args=[
                            [None],
                            dict(
                                frame=dict(duration=0, redraw=False),
                                mode="immediate",
                                transition=dict(duration=0),
                            ),
                        ],
                    ),
                ],
                x=0.5,
                y=-0.2,
            )
        ],
    )

    return fig, mse


def Plot_featue_importance(test_split):
    try:
        feature_importance = reg.feature_importances_
    except:
        # return blank figures
        fig = go.Figure()
        fig.update_layout(title="Train a model to see the plots")
        return fig, fig

    sorted_idx = np.argsort(feature_importance)

    fig = px.bar(
        pd.DataFrame(
            {
                "Importance": feature_importance[sorted_idx],
                "Feature": np.array(diabetes.feature_names)[sorted_idx],
            }
        ),
        x="Importance",
        y="Feature",
        orientation="h",
        title="Feature Importance (MDI)",
    )

    X, y = diabetes.data, diabetes.target
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_split, random_state=42
    )
    result = permutation_importance(
        reg, X_test, y_test, n_repeats=10, random_state=42, n_jobs=2
    )

    fig1 = px.box(
        pd.DataFrame(result.importances.T, columns=diabetes.feature_names),
        title="Permutation Importance (test set)",
    )

    return fig, fig1


with gr.Blocks() as demo:
    gr.Markdown("# Gradient Boosting regression")
    gr.Markdown(
        "This demo is based on [gradient boosting regression example of scikit-learn](https://scikit-learn.org/stable/auto_examples/ensemble/plot_gradient_boosting_regression.html) Example.This example demonstrates gradient goosting to produce a predictive model from an ensemble of weak predictive models. Gradient boosting can be used for regression and classification problems. Here, we will train a model to tackle a diabetes regression task."
    )

    with gr.Tab("Train the model"):
        gr.Markdown("### Below is the diabetes dataset used in this demo 👇 ")
        gr.Markdown("### You can change the interval of rows to display.")
        gr.Markdown(
            "The diabetes dataset consists of ten baseline variables, age, sex, body mass index (BMI), average blood pressure (BP), and six blood serum measurements for 442 diabetes patients. The target variable is a quantitative measure of disease progression one year after baseline."
        )
        total_rows = X.shape[0]
        rows_number = gr.Slider(
            0, total_rows, label="Displaying Rows", value=5, step=5
        )

        rows_number.change(
            fn=display_table,
            inputs=[rows_number],
            outputs=[gr.Text(label="Row"), gr.DataFrame()],
        )

        gr.Markdown(
            "# Play with the parameters to see how the model performance changes"
        )

        gr.Markdown(
            """
        ### `Number of Estimators` : the number of boosting stages that will be performed. Later, we will plot deviance against boosting iterations.

        ### `Max Depth` : limits the number of nodes in the tree. The best value depends on the interaction of the input variables.

        ### `Min Samples Split` : the minimum number of samples required to split an internal node.

        ### `learning_rate` : how much the contribution of each tree will shrink.

        ### `loss` : loss function to optimize.

        ### `Test Split` : the percentage of the dataset to include in the test split.

        ### `Animation Speed for Deviance Plot` : the duration of the animation of Deviation Plot.
        """
        )
        with gr.Row():
            test_split = gr.Slider(0.1, 0.9, label="Test Split", value=0.2, step=0.1)
            learning_rate = gr.Slider(
                0.01, 0.5, label="Learning Rate", value=0.1, step=0.01
            )
            n_estimators = gr.Slider(
                10, 1000, label="Number of Estimators", value=100, step=10
            )
            max_depth = gr.Slider(1, 10, label="Max Depth", value=3, step=1)
            min_samples_split = gr.Slider(
                2, 10, label="Min Samples Split", value=2, step=1
            )
            loss = gr.Dropdown(
                ["squared_error", "absolute_error", "huber", "quantile"],
                label="Loss",
                value="squared_error",
            )

            duration = gr.Slider(
                0, 100, label="Animation Speed for Deviance Plot", value=25, step=10
            )

        model_btn = gr.Button("Train Model")
        gr.Markdown(
            "### Finally, we will visualize the results. To do that, we will first compute the test set deviance and then plot it against boosting iterations."
        )
        model_btn.click(
            fn=train_model,
            inputs=[
                test_split,
                learning_rate,
                n_estimators,
                max_depth,
                min_samples_split,
                loss,
                duration,
            ],
            outputs=[gr.Plot(), gr.Text(label="MSE")],
        )

    with gr.Tab("Feature Importance"):
        gr.Markdown("## Feature Importance (MDI) and Permutation Importance (test set)")
        gr.Markdown(
            "For this example, the impurity-based and permutation methods identify the same 2 strongly predictive features but not in the same order. The third most predictive feature, “bp”, is also the same for the 2 methods. The remaining features are less predictive and the error bars of the permutation plot show that they overlap with 0."
        )
        feat_imp_btn = gr.Button("Plot Feature Importance")
        with gr.Row():
            feat_imp_btn.click(
                fn=Plot_featue_importance,
                inputs=[test_split],
                outputs=[gr.Plot(), gr.Plot()],
            )


demo.launch()