# Import libraries
import yfinance as yf
import tensorflow as tf
import numpy as np
import gradio as gr
from IPython.display import display

# Set up TensorBoard log directory
log_dir = 'tmp'

tf.config.experimental.set_virtual_device_configuration = None
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if physical_devices:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

# Define the function to predict asset prices
def predict_asset(ticker):
    # Download a year's worth of data for the asset the user entered
    data = yf.download(ticker, period='5y')

    # Normalize the data using TensorFlow's Keras utility
    matrix = tf.keras.utils.normalize(data.values)

    # Create a sequential machine learning model using TensorFlow's Keras library
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(128, activation='relu', input_shape=(matrix.shape[1],)))
    model.add(tf.keras.layers.Dense(64, activation='relu'))
    model.add(tf.keras.layers.Dense(1))
    model.compile(optimizer='adam', loss='mean_squared_error')

    # Create TensorBoard callback
    tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir)

    # Train the model using the normalized data and the asset prices
    model.fit(matrix, data['Close'], epochs=100, callbacks=[tensorboard_callback])

    # Use the trained model to predict prices for the asset
    predicted_prices = model.predict(matrix)
    predicted_prices = np.array(predicted_prices).flatten()

    # Calculate the average predicted price
    mean_predicted_price = np.mean(predicted_prices)
    current_price = data['Close'][-1]
    percent_change = (mean_predicted_price - current_price) / current_price * 100

    # Set variables to control the color and wording of the output
    color = 'green' if percent_change >= 0 else 'red'
    increase_decrease = 'increase' if percent_change >= 0 else 'decrease'

    # Return the formatted output as a string
    return '<h1>' + ticker + ' Predicted Prices 5 years in the future:</h1><br />' + str(predicted_prices) + \
           '<br /><br /><h1>' + ticker + ' Average Predicted Price 5 years in the Future:<br />&dollar;' + str(mean_predicted_price) + \
           '<br /><br />' + ticker + ' Current Price:<br />&dollar;' + str(current_price) + \
           '<br /><br />' + ticker + ' Percent Change from Current Price to Average Predicted Price 5 Years in the Future:<br /><span style="color: ' + color + '">' + \
           str(percent_change) + '% ' + increase_decrease + '</span></h1>'

# Create the Gradio interface
interface = gr.Interface(
    fn=predict_asset,
    inputs="text",
    outputs="html",
    title="Asset Price Prediction",
    description="Enter a ticker to predict asset price 5 years in the future.",
)

# Launch the interface
interface.launch(share=True, debug=True, auth=("admin", "password12345"))