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

def load_and_scale_data(file_path, number_of_columns):
    """Loads and scales stock data using MinMaxScaler."""
    data = pd.read_csv(file_path, header=None, delimiter='\t')
    scalers = {}
    for j in range(number_of_columns):
        scalers[j] = MinMaxScaler()
        data.iloc[:, j] = scalers[j].fit_transform(data.iloc[:, j].values.reshape(-1, 1)).ravel()
    return data, scalers

def prepare_datasets(data, look_back, look_forward):
    """Prepares training datasets."""
    features, targets = [], []
    for index in range(look_back, len(data) - look_forward + 1):
        features.append(data.iloc[index - look_back:index, :].values)
        targets.append(data.iloc[index + look_forward - 1, :].values)
    return np.array(features), np.array(targets)

def train_model(model, training_loader, optimizer, loss_func, num_epochs):
    """Trains the LSTM model."""
    for epoch_num in range(num_epochs):
        model.train()
        for X_batch, y_batch in training_loader:
            optimizer.zero_grad()
            predictions = model(X_batch)
            loss = loss_func(predictions, y_batch)
            loss.backward()
            optimizer.step()
        if (epoch_num + 1) % 10 == 0:
            print(f'Epoch {epoch_num + 1}/{num_epochs}, Loss: {loss.item()}')
        if (epoch_num + 1) % 100 == 0:
            torch.save(model.state_dict(), str(epoch_num)+'_model.pth')
            
def main():
    # Constants and Parameters
    dataset_path = 'AI_Lec_23_final_stocks.csv'
    look_back = 20
    total_stocks = 100
    num_epochs = 500
    alpha = 0.0001

    # Load and Scale Data
    stock_data, scalers = load_and_scale_data(dataset_path, total_stocks)

    # Prepare Datasets
    X, y = prepare_datasets(stock_data, look_back, 1)
    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_tensor = torch.tensor(X_train, dtype=torch.float32)
    y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
    training_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    training_loader = DataLoader(training_dataset, batch_size=16, shuffle=True)

    # Initialize Model
    deep_stock_model = DeepLSTM(input_size=total_stocks, hidden_size=128, layers_count=3, output_size=total_stocks, dropout_value=0.3)
    optimizer = torch.optim.Adam(deep_stock_model.parameters(), lr=alpha)
    loss_function = nn.MSELoss()

    # Train the Model
    train_model(deep_stock_model, training_loader, optimizer, loss_function, num_epochs)

    # Save the Trained Model
    torch.save(deep_stock_model.state_dict(), 'deep_lstm_stock_predictor_model.pth')

if __name__ == "__main__":
    main()