import streamlit as st
import numpy as np
import pandas as pd
import tensorflow as tf
from matplotlib import pyplot as plt

# Function to build the model
def build_model(my_learning_rate):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(units=1, input_shape=(1,)))
    model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=my_learning_rate),
                  loss='mean_squared_error',
                  metrics=[tf.keras.metrics.RootMeanSquaredError()])
    return model

# Function to train the model
def train_model(model, feature, label, epochs, batch_size):
    history = model.fit(x=feature, y=label, batch_size=batch_size, epochs=epochs)
    trained_weight = model.get_weights()[0][0]
    trained_bias = model.get_weights()[1]
    epochs = history.epoch
    hist = pd.DataFrame(history.history)
    rmse = hist["root_mean_squared_error"]
    return trained_weight, trained_bias, epochs, rmse

# Function to plot the model
def plot_the_model(trained_weight, trained_bias, feature, label):
    plt.figure(figsize=(10, 6))
    plt.xlabel('Feature')
    plt.ylabel('Label')

    # Plot the feature values vs. label values
    plt.scatter(feature, label, c='b')

    # Create a red line representing the model
    x0 = 0
    y0 = trained_bias
    x1 = feature[-1][0]
    y1 = trained_bias + (trained_weight * x1)
    plt.plot([x0, x1], [y0, y1], c='r')

    st.pyplot(plt)

# Function to plot the loss curve
def plot_the_loss_curve(epochs, rmse):
    plt.figure(figsize=(10, 6))
    plt.xlabel('Epoch')
    plt.ylabel('Root Mean Squared Error')

    plt.plot(epochs, rmse, label='Loss')
    plt.legend()
    plt.ylim([rmse.min()*0.97, rmse.max()])
    st.pyplot(plt)

# Define the dataset
my_feature = np.array([1.0, 2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0,  9.0, 10.0, 11.0, 12.0], dtype=float).reshape(-1, 1)
my_label   = np.array([5.0, 8.8,  9.6, 14.2, 18.8, 19.5, 21.4, 26.8, 28.9, 32.0, 33.8, 38.2], dtype=float)

# Streamlit interface
st.title("Simple Linear Regression with Synthetic Data")

st.write("https://colab.research.google.com/github/google/eng-edu/blob/main/ml/cc/exercises/linear_regression_with_synthetic_data.ipynb?utm_source=mlcc&utm_campaign=colab-external&utm_medium=referral&utm_content=linear_regression_synthetic_tf2-colab&hl=en")

learning_rate = st.sidebar.slider('Learning Rate', min_value=0.001, max_value=1.0, value=0.01, step=0.01)
epochs = st.sidebar.slider('Epochs', min_value=1, max_value=1000, value=10, step=1)
batch_size = st.sidebar.slider('Batch Size', min_value=1, max_value=len(my_feature), value=12, step=1)

if st.sidebar.button('Run'):
    my_model = build_model(learning_rate)
    trained_weight, trained_bias, epochs, rmse = train_model(my_model, my_feature, my_label, epochs, batch_size)

    st.subheader('Model Plot')
    plot_the_model(trained_weight, trained_bias, my_feature, my_label)

    st.subheader('Loss Curve')
    plot_the_loss_curve(epochs, rmse)