import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import io

def create_features(data, look_back=1):
    X, y = [], []
    for i in range(len(data) - look_back):
        X.append(data[i:(i + look_back)])
        y.append(data[i + look_back])
    return np.array(X), np.array(y)

def train_and_predict(file):
    # 載入和預處理數據
    dataframe = pd.read_csv(file, usecols=[1], engine='python')
    dataset = dataframe.values.astype('float32')
    
    # 正規化數據集
    scaler = MinMaxScaler(feature_range=(0, 1))
    dataset = scaler.fit_transform(dataset)
    
    # 創建特徵和目標
    look_back = 3
    X, y = create_features(dataset, look_back)
    
    # 分割訓練集和測試集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
    
    # 創建和訓練模型
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    # 進行預測
    train_predict = model.predict(X_train)
    test_predict = model.predict(X_test)
    
    # 反轉預測
    train_predict = scaler.inverse_transform(train_predict.reshape(-1, 1))
    y_train = scaler.inverse_transform(y_train.reshape(-1, 1))
    test_predict = scaler.inverse_transform(test_predict.reshape(-1, 1))
    y_test = scaler.inverse_transform(y_test.reshape(-1, 1))
    
    # 計算 RMSE
    train_score = np.sqrt(mean_squared_error(y_train, train_predict))
    test_score = np.sqrt(mean_squared_error(y_test, test_predict))
    
    # 準備繪圖數據
    train_predict_plot = np.empty_like(dataset)
    train_predict_plot[:, :] = np.nan
    train_predict_plot[look_back:len(train_predict)+look_back, :] = train_predict
    test_predict_plot = np.empty_like(dataset)
    test_predict_plot[:, :] = np.nan
    test_predict_plot[len(train_predict)+(look_back*2)+1:len(dataset)-1, :] = test_predict
    
    return scaler.inverse_transform(dataset), train_predict_plot, test_predict_plot, train_score, test_score

st.title("Stock Price Prediction")

uploaded_file = st.file_uploader("Choose a CSV file", type="csv")

if uploaded_file is not None:
    # 讀取文件並進行預測
    original_data, train_predict, test_predict, train_score, test_score = train_and_predict(uploaded_file)
    
    # 顯示結果
    st.subheader("Prediction Results")
    
    # 繪製圖表
    fig, ax = plt.subplots(figsize=(12, 6))
    ax.plot(original_data, label='Original Data', color='blue')
    ax.plot(train_predict, label='Training Predictions', linestyle='--', color='green')
    ax.plot(test_predict, label='Test Predictions', linestyle='--', color='red')
    ax.set_xlabel('Time')
    ax.set_ylabel('Stock Price')
    ax.set_title('Original Data and Predictions')
    ax.legend()
    ax.grid(True, linestyle='--', alpha=0.7)
    st.pyplot(fig)
    
    # 顯示評分
    col1, col2 = st.columns(2)
    with col1:
        st.metric("Train Score (RMSE)", f"{train_score:.2f}")
    with col2:
        st.metric("Test Score (RMSE)", f"{test_score:.2f}")

else:
    st.info("Please upload a CSV file to start the prediction.")

st.markdown("""
This application uses a simple linear regression model to predict stock prices based on historical data.

To use:
1. Upload a CSV file containing historical stock price data.
2. The app will train a model on this data and make predictions.
3. You'll see a graph showing the original data and predictions, along with RMSE scores for training and test sets.

Note: The CSV file should have the stock prices in the second column.
""")