# Gradio Implementation: Lenix Carter
# License: BSD 3-Clause or CC-0

import gradio as gr

import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

plt.switch_backend("agg")

def compare_reg(n_samples, n_features):
    np.random.seed(42)

    X = np.random.randn(n_samples, n_features)
    true_coef = 3 * np.random.randn(n_features)

    # Threshold coefficients to render them non-negative
    true_coef[true_coef < 0] = 0
    y = np.dot(X, true_coef)

    # Add some noise
    y += 5 * np.random.normal(size=(n_samples,))

    # Split the data in train set and test set
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5)

    reg_nnls = LinearRegression(positive=True)
    y_pred_nnls = reg_nnls.fit(X_train, y_train).predict(X_test)
    r2_score_nnls = r2_score(y_test, y_pred_nnls)

    reg_ols = LinearRegression()
    y_pred_ols = reg_ols.fit(X_train, y_train).predict(X_test)
    r2_score_ols = r2_score(y_test, y_pred_ols)

    fig, ax = plt.subplots()
    ax.plot(reg_ols.coef_, reg_nnls.coef_, linewidth=0, marker=".")

    low_x, high_x = ax.get_xlim()
    low_y, high_y = ax.get_ylim()
    low = max(low_x, low_y)
    high = min(high_x, high_y)
    ax.plot([low, high], [low, high], ls="--", c=".3", alpha=0.5)
    ax.set_xlabel("OLS regression coefficients", fontweight="bold")
    ax.set_ylabel("NNLS regression coefficients", fontweight="bold")

    scores = "The R2 for NNLS is {}\nThe R2 for OLS is {}".format(r2_score_nnls, r2_score_ols)
    return fig, scores

title = "Non-negative Least Squares"
with gr.Blocks() as demo:
    gr.Markdown(f" # {title}")
    gr.Markdown("""
                     This example fits a linear model with positivity constraints on the regression coefficients and compares the estimated coefficients to a classic linear regression. 

                     This is based on the example [here](https://scikit-learn.org/stable/auto_examples/linear_model/plot_nnls.html#sphx-glr-auto-examples-linear-model-plot-nnls-py).
                     """)
    with gr.Row():
        with gr.Column():
            n_samp = gr.Slider(100, 1000, 200, step=1, label="Number of Samples")
            n_feat = gr.Slider(3, 100, 50, step=1, label="Number of Features")
        with gr.Column():
            scores = gr.Textbox(label="R2 Scores")
            coeff_comp_graph = gr.Plot(label="Comparison of Coefficients")
    n_samp.change(
                  fn=compare_reg,
                  inputs=[n_samp, n_feat],
                  outputs=[coeff_comp_graph, scores]
                  )
    n_feat.change(
                  fn=compare_reg,
                  inputs=[n_samp, n_feat],
                  outputs=[coeff_comp_graph, scores]
                  )
    with gr.Row():
        gr.Markdown("This shows a high degree of correlation between the the regression coefficients of OLS and NNLS. However, we observe that some coefficients in the NNLS regression shrink to 0.")

if __name__ == '__main__':
    demo.launch()