import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

import matplotlib.cm as cm
from sklearn.utils import shuffle
from sklearn.utils import check_random_state
from sklearn.linear_model import BayesianRidge

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

description = f"""
## Description

This demo computes a Bayesian Ridge Regression of Sinusoids.

In general, when fitting a curve with a polynomial by Bayesian ridge regression, the selection of initial values of the regularization parameters (alpha, lambda) may be important. This is because the regularization parameters are determined by an iterative procedure that depends on initial values.

In this demo, the sinusoid is approximated by a cubic polynomial using different pairs of initial values.  By evaluating log marginal likelihood (L) of these models, we can determine which one is better. It can be concluded that the model with larger L is more likely.

The demo is based on the [scikit-learn docs](https://scikit-learn.org/stable/auto_examples/linear_model/plot_bayesian_ridge_curvefit.html#sphx-glr-auto-examples-linear-model-plot-bayesian-ridge-curvefit-py)
"""

def func(x):
    return np.sin(2 * np.pi * x)

def curve_fit(size, alpha, lam):
    rng = np.random.RandomState(1234)
    x_train = rng.uniform(0.0, 1.0, size)
    y_train = func(x_train) + rng.normal(scale=0.1, size=size)
    x_test = np.linspace(0.0, 1.0, 100)

    n_order = 3
    X_train = np.vander(x_train, n_order + 1, increasing=True)
    X_test = np.vander(x_test, n_order + 1, increasing=True)
    reg = BayesianRidge(tol=1e-6, fit_intercept=False, compute_score=True)

    fig1, ax1 = plt.subplots(1, 1, figsize=(8, 4))
    fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4))

    # Bayesian ridge regression with different initial value pairs
    init = [1 / np.var(y_train), 1.0]  # Default values
    reg.fit(X_train, y_train)
    ymean, ystd = reg.predict(X_test, return_std=True)
    
    ax1.plot(x_test, func(x_test), color="blue", label="sin($2\\pi x$)")
    ax1.scatter(x_train, y_train, s=50, alpha=0.5, label="observation")
    ax1.plot(x_test, ymean, color="red", label="predict mean")
    ax1.fill_between(
        x_test, ymean - ystd, ymean + ystd, color="pink", alpha=0.5, label="predict std"
    )
    ax1.set_ylim(-1.3, 1.3)
    ax1.legend()
    title = "$\\alpha$_init$={:.2f},\\ \\lambda$_init$={}$".format(init[0], init[1])
    title += " (Default)"
    ax1.set_title(title, fontsize=12)
    text = "$\\alpha={:.1f}$\n$\\lambda={:.3f}$\n$L={:.1f}$".format(
        reg.alpha_, reg.lambda_, reg.scores_[-1]
    )
    ax1.text(0.05, -1.0, text, fontsize=12)

    init = [alpha, lam]
    reg.set_params(alpha_init=init[0], lambda_init=init[1])
    reg.fit(X_train, y_train)
    ymean, ystd = reg.predict(X_test, return_std=True)

    ax2.plot(x_test, func(x_test), color="blue", label="sin($2\\pi x$)")
    ax2.scatter(x_train, y_train, s=50, alpha=0.5, label="observation")
    ax2.plot(x_test, ymean, color="red", label="predict mean")
    ax2.fill_between(
        x_test, ymean - ystd, ymean + ystd, color="pink", alpha=0.5, label="predict std"
    )
    ax2.set_ylim(-1.3, 1.3)
    ax2.legend()
    title = "$\\alpha$_init$={:.2f},\\ \\lambda$_init$={}$".format(init[0], init[1])
    
    title += " (Chosen)"
    ax2.set_title(title, fontsize=12)
    text = "$\\alpha={:.1f}$\n$\\lambda={:.3f}$\n$L={:.1f}$".format(
        reg.alpha_, reg.lambda_, reg.scores_[-1]
    )
    ax2.text(0.05, -1.0, text, fontsize=12)
    
    return fig1, fig2


with gr.Blocks(theme=theme) as demo:
    gr.Markdown('''
            <h1 style='text-align: center'>Curve Fitting with Bayesian Ridge Regression 📈</h1>
        ''')
    gr.Markdown(description)
    
    with gr.Row():
        size = gr.Slider(minimum=10, maximum=100, step=5, value=25, label="Number of Data Points")
        alpha = gr.Slider(minimum=1e-2, maximum=2, step=0.1, value=1, label="Initial Alpha")
        lam = gr.Slider(minimum=1e-5, maximum=1, step=1e-4, value=1e-3, label="Initial Lambda")
    with gr.Row():
        run_button = gr.Button('Fit the Curve')
    with gr.Row():
        plot_default = gr.Plot(label = 'Default Initialization')
        plot_chosen = gr.Plot(label = 'Chosen Initialization')
    run_button.click(fn=curve_fit, inputs=[size, alpha, lam], outputs=[plot_default, plot_chosen])

demo.launch()