import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import mean_squared_error

st.subheader("Ridge Demo")
col1, col2 = st.columns(2)

degree = st.slider('Degree', 2, 40, 1)
alpha = st.slider('Lambda (Regularisation)', 0, 500, 1)


with col1:
    st.markdown("#### Un-regularized")
    
with col2:
    st.markdown("#### Regularized")
    
x = np.linspace(-1., 1., 100)
y = 4 + 3*x + 2*np.sin(x) + 2*np.random.randn(len(x))


poly = PolynomialFeatures(degree=degree, include_bias=False)
x_new = poly.fit_transform(x.reshape(-1, 1))

lr = LinearRegression()
lr.fit(x_new, y)
y_pred = lr.predict(x_new)


ri = Ridge(alpha = alpha)
ri.fit(x_new, y)
y_pred_ri = ri.predict(x_new)


fig1, ax1 = plt.subplots()
fig2, ax2 = plt.subplots()

ax1.scatter(x, y)
ax1.plot(x, y_pred)

ax2.scatter(x, y)
ax2.plot(x, y_pred_ri)

for ax in [ax1, ax2]:
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    # Only show ticks on the left and bottom spines
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')

    ax.set_xlabel("x")
    ax.set_ylabel("y")

rmse = np.round(np.sqrt(mean_squared_error(y_pred, y)), 2)
ax1.set_title(f"Train RMSE: {rmse}")

rmse_ri = np.round(np.sqrt(mean_squared_error(y_pred_ri, y)), 2)
ax2.set_title(f"Train RMSE: {rmse_ri}")

with col1:
    st.pyplot(fig1)
    
with col2:
    st.pyplot(fig2)
hide_streamlit_style = """
            <style>
            #MainMenu {visibility: hidden;}
            footer {visibility: hidden;}
            </style>
            """
st.markdown(hide_streamlit_style, unsafe_allow_html=True)