import gradio as gr
import numpy as np
import matplotlib.pyplot as plt

from sklearn.metrics import r2_score
from sklearn.linear_model import Lasso, ElasticNet


theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
)
model_card = f"""
## Description

This demo estimates **Lasso** and **Elastic-Net** regression models on a manually generated sparse signal corrupted with an additive noise.
You can play around with different ``regularization strength``, ``mixing ratio between L1 and L2``, ``number of samples``, ``number of features`` to see the effect

## Dataset

Simulation dataset
"""



def do_train(alpha, l1_ratio, n_samples, n_features):
    np.random.seed(42)
    X = np.random.randn(n_samples, n_features)

    # Decreasing coef w. alternated signs for visualization
    idx = np.arange(n_features)
    coef = (-1) ** idx * np.exp(-idx / 10)
    coef[10:] = 0  # sparsify coef
    y = np.dot(X, coef)

    # Add noise
    y += 0.01 * np.random.normal(size=n_samples)

    # Split data in train set and test set
    n_samples = X.shape[0]
    X_train, y_train = X[: n_samples // 2], y[: n_samples // 2]
    X_test, y_test = X[n_samples // 2 :], y[n_samples // 2 :]

    lasso = Lasso(alpha=alpha)
    y_pred_lasso = lasso.fit(X_train, y_train).predict(X_test)
    r2_score_lasso = r2_score(y_test, y_pred_lasso)


    enet = ElasticNet(alpha=alpha, l1_ratio=l1_ratio)
    y_pred_enet = enet.fit(X_train, y_train).predict(X_test)
    r2_score_enet = r2_score(y_test, y_pred_enet)

    fig, axes = plt.subplots()

    m, s, _ = axes.stem(
    np.where(enet.coef_)[0],
    enet.coef_[enet.coef_ != 0],
    markerfmt="x",
    label="Elastic net coefficients",
    )
    plt.setp([m, s], color="#2ca02c")

    m, s, _ = plt.stem(
        np.where(lasso.coef_)[0],
        lasso.coef_[lasso.coef_ != 0],
        markerfmt="x",
        label="Lasso coefficients",
    )
    plt.setp([m, s], color="#ff7f0e")

    axes.stem(
        np.where(coef)[0],
        coef[coef != 0],
        label="True coefficients",
        markerfmt="bx",
    )

    axes.legend(loc="best")
    axes.set_title("Elastic net and Lasso coefficients")
    text = f"Lasso R^2: {r2_score_lasso:.3f}, Elastic Net R^2: {r2_score_enet:.3f}"
    return fig, text



with gr.Blocks(theme=theme) as demo:
    gr.Markdown('''
            <div>
            <h1 style='text-align: center'>Lasso and Elastic Net for Sparse Signals</h1>
            </div>
        ''')
    gr.Markdown(model_card)
    gr.Markdown("Author: <a href=\"https://huggingface.co/vumichien\">Vu Minh Chien</a>. Based on the example from <a href=\"https://scikit-learn.org/stable/auto_examples/linear_model/plot_lasso_and_elasticnet.html#sphx-glr-auto-examples-linear-model-plot-lasso-and-elasticnet-py\">scikit-learn</a>")
    alpha = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.1, label="Controlling regularization strength: alpha. Using alpha = 0 with the Lasso object is not advised")
    l1_ratio = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label="The ElasticNet mixing parameter: l1_ratio. For l1_ratio = 0 the penalty is an L2 penalty. For l1_ratio = 1 it is an L1 penalty. For 0 < l1_ratio < 1, the penalty is a combination of L1 and L2.")
    n_samples = gr.Slider(minimum=50, maximum=500, step=50, value=50, label="Number of samples")
    n_features = gr.Slider(minimum=50, maximum=200, step=50, value=50, label="Number of features")
    with gr.Row():
        with gr.Column():
            plot = gr.Plot(label="Coefficients plot")
        with gr.Column():
            results = gr.Textbox(label="Results")

    alpha.change(fn=do_train, inputs=[alpha, l1_ratio, n_samples, n_features], outputs=[plot, results])
    l1_ratio.change(fn=do_train, inputs=[alpha, l1_ratio, n_samples, n_features], outputs=[plot, results])
    n_samples.change(fn=do_train, inputs=[alpha, l1_ratio, n_samples, n_features], outputs=[plot, results])
    n_features.change(fn=do_train, inputs=[alpha, l1_ratio, n_samples, n_features], outputs=[plot, results])

demo.launch()