import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import StandardScaler

# Define the LSTM model with Attention
class Attention(nn.Module):
    def __init__(self, feature_dim, step_dim, bias=True, **kwargs):
        super(Attention, self).__init__(**kwargs)
        self.feature_dim = feature_dim
        self.step_dim = step_dim
        self.bias = bias
        self.attention_weights = nn.Parameter(torch.Tensor(feature_dim, 1))
        nn.init.kaiming_uniform_(self.attention_weights, a=0, mode='fan_in', nonlinearity='leaky_relu')
        self.bias = nn.Parameter(torch.Tensor(1)) if bias else None

    def forward(self, x):
        eij = torch.mm(x.contiguous().view(-1, self.feature_dim), self.attention_weights).view(-1, self.step_dim)
        eij = eij + self.bias if self.bias is not None else eij
        eij = torch.tanh(eij)
        a = torch.exp(eij)
        a = a / (torch.sum(a, 1, keepdim=True) + 1e-10)
        weighted_input = x * torch.unsqueeze(a, -1)
        return torch.sum(weighted_input, 1)

class LSTMWithAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim, dropout_prob):
        super(LSTMWithAttention, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, dropout=dropout_prob)
        self.attention = Attention(hidden_dim, n_past)
        self.fc = nn.Linear(hidden_dim, output_dim * n_future)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
        out, _ = self.lstm(x, (h0.detach(), c0.detach()))
        out = self.attention(out)
        out = self.fc(out)
        return out

# Constants
n_past = 20
n_future = 1
n_features = 100

# Load the trained model
model_path = 'final_model.pth'  # Update with your model file path
model = LSTMWithAttention(input_dim=n_features, hidden_dim=100, num_layers=2, output_dim=n_features, dropout_prob=0.2)
model.load_state_dict(torch.load(model_path))
model.eval()

# Load your dataset
file_path = 'Final_stocks.csv'  # Replace with your file path
data = pd.read_csv(file_path, header=None, sep='\t')

# Initialize and apply StandardScaler for each stock
scalers = {}
for i in range(data.shape[1]):
    scalers[i] = StandardScaler()
    data.iloc[:, i] = scalers[i].fit_transform(data.iloc[:, i].values.reshape(-1, 1)).flatten()

# Predicting future 20 days
predicted_data = data.copy()
for day in range(20):
    last_20_days = np.array([predicted_data.iloc[-n_past:, :]])
    last_20_days_tensor = torch.from_numpy(last_20_days).float()

    # Predicting the next day
    with torch.no_grad():
        next_day_prediction = model(last_20_days_tensor)
        next_day_prediction = next_day_prediction.view(-1, n_features).numpy()

    # Inverse transform the prediction to original scale
    for i in range(n_features):
        next_day_prediction[:, i] = scalers[i].inverse_transform(next_day_prediction[:, i].reshape(-1, 1)).flatten()

    # Create a DataFrame for the next day prediction
    next_day_df = pd.DataFrame(next_day_prediction, columns=predicted_data.columns)

    # Append the prediction for the next day using concat
    predicted_data = pd.concat([predicted_data, next_day_df], ignore_index=True)

# Save the last 20 days of predictions to a CSV file
predicted_data.tail(20).to_csv('predicted_stock_prices.csv', header=False, sep='\t', index=False)
print("Saved the predicted stock prices for the next 20 days.")