import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import MinMaxScaler
from model import DeepLSTM

def load_model(model_path, input_size, hidden_size, layers_count, output_size, dropout_value):
    """Loads the trained deep LSTM model."""
    model = DeepLSTM(input_size=input_size, hidden_size=hidden_size, layers_count=layers_count, output_size=output_size, dropout_value=dropout_value)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    return model

def scale_data(stock_data):
    """Applies MinMaxScaler to the stock data."""
    scalers = {}
    for i in range(stock_data.shape[1]):
        scalers[i] = MinMaxScaler()
        stock_data.iloc[:, i] = scalers[i].fit_transform(stock_data.iloc[:, i].values.reshape(-1, 1)).ravel()
    return stock_data, scalers

def predict_future_prices(model, data, scalers, historical_days, future_days):
    """Predicts future stock prices."""
    future_prediction_data = data.copy()
    for day in range(future_days):
        recent_data = np.array([future_prediction_data.iloc[-historical_days:, :]])
        recent_data_tensor = torch.tensor(recent_data, dtype=torch.float32)

        with torch.no_grad():
            next_day_forecast = model(recent_data_tensor)
            next_day_forecast = next_day_forecast.view(-1, data.shape[1]).numpy()

        for i in range(data.shape[1]):
            next_day_forecast[:, i] = scalers[i].inverse_transform(next_day_forecast[:, i].reshape(-1, 1)).ravel()

        next_day_df = pd.DataFrame(next_day_forecast, columns=future_prediction_data.columns)
        future_prediction_data = pd.concat([future_prediction_data, next_day_df], ignore_index=True)

    return future_prediction_data.tail(future_days)

def main():
    # Constants
    historical_days = 20
    number_of_stocks = 100
    future_days = 20
    model_path = 'deep_lstm_stock_predictor_model.pth'
    data_path = 'AI_Lec_23_final_stocks.csv'

    # Load the model
    model = load_model(model_path, number_of_stocks, 128, 3, number_of_stocks, 0.3)

    # Load and scale the data
    stock_data = pd.read_csv(data_path, header=None, sep='\t')
    scaled_data, scalers = scale_data(stock_data)

    # Predict future prices
    future_prices = predict_future_prices(model, scaled_data, scalers, historical_days, future_days)

    # Save the predictions
    future_prices.to_csv('predicted_prices.csv', header=False, sep='\t', index=False)
    print("Predicted future stock prices are saved.")

if __name__ == "__main__":
    main()