# A simple Linear Regression example with TensorFlow

import tensorflow as tf
import numpy as np
import streamlit as st
import matplotlib.pyplot as plt

# Define the model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(units=1, input_shape=[1])
])

# Compile the model with an optimizer and loss function
model.compile(optimizer='sgd', loss='mse')

# Training data
xs = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=float)
ys = np.array([1.5, 2.0, 2.5, 3.0, 3.5], dtype=float)

# Streamlit UI
st.title('Simple Linear Regression with TensorFlow')

# User input for the new value to predict
input_value = st.number_input('Enter your input value:', value=1.0, format="%.1f")

# User input for epochs
epochs = st.sidebar.slider("Number of epochs", 10, 100, 10)

# Button to train the model and make prediction
if st.button('Train Model and Predict'):
    with st.spinner('Training...'):
        model.fit(xs, ys, epochs=epochs)
    st.success('Training completed!')

    # Make prediction
    prediction = model.predict([input_value])
    st.write(f'For input {input_value}, the prediction is {prediction[0][0]}')

    # Predictions for visualization
    predictions = model.predict(xs)

    # Plotting
    plt.scatter(xs, ys, label='Actual')
    plt.plot(xs, predictions, color='red', label='Predicted')
    plt.xlabel('Input Feature')
    plt.ylabel('Output Value')
    plt.legend()
    st.pyplot(plt)