import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader

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

# Constants
n_past = 20
n_future = 1
n_features = 100  # For 100 stocks
epochs = 100
learning_rate = 0.0002

# Initialize 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()

# Prepare the training data
X, y = [], []
for i in range(n_past, len(data) - n_future):
    X.append(data.iloc[i - n_past:i, :].values)
    y.append(data.iloc[i + n_future - 1, :].values)

X, y = np.array(X), np.array(y)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Convert to PyTorch tensors
X_train = torch.from_numpy(X_train).float()
y_train = torch.from_numpy(y_train).float()
X_test = torch.from_numpy(X_test).float()
y_test = torch.from_numpy(y_test).float()

# TensorDataset and DataLoader
train_data = TensorDataset(X_train, y_train)
batch_size = 16
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)

# Attention mechanism
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)

# LSTM with Attention
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

# Model, Optimizer, and Loss
model = LSTMWithAttention(input_dim=n_features, hidden_dim=100, num_layers=2, output_dim=n_features, dropout_prob=0.2)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.MSELoss()

# Training Loop with Test Loss Calculation
for epoch in range(epochs):
    model.train()
    train_loss = 0.0
    for X_batch, y_batch in train_loader:
        optimizer.zero_grad()
        output = model(X_batch)
        loss = criterion(output, y_batch)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    # Calculate test loss
    model.eval()
    test_loss = 0.0
    with torch.no_grad():
        for X_batch, y_batch in DataLoader(TensorDataset(X_test, y_test), batch_size=batch_size):
            output = model(X_batch)
            loss = criterion(output, y_batch)
            test_loss += loss.item()

    # Average losses
    train_loss /= len(train_loader)
    test_loss /= len(X_test) // batch_size

    print(f'Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')


# save the final model after all epochs
torch.save(model.state_dict(), 'final_model.pth')

# 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.")