# %% Import required packages
import torch
from src.models.transformer_model import TransformerModel
from src.models.rl_model import RLModel
from src.models.trading_agent import TradingAgent
from src.training.train import train_transformer, train_rl
from src.data.data_preprocessing import load_processed_data

# %% Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# %% Load processed data
data = load_processed_data('./data/processed/processed_data.csv')

# %% Initialize models
transformer_model = TransformerModel().to(device)
rl_model = RLModel().to(device)
trading_agent = TradingAgent(transformer_model, rl_model)

# %% Train Transformer Model
# Set the appropriate hyperparameters
transformer_model_hyperparams = {
    "epochs": 10,
    "batch_size": 32,
    "learning_rate": 0.001,
}
train_transformer(transformer_model, data, transformer_model_hyperparams)

# %% Save Transformer Model
torch.save(transformer_model.state_dict(), './models/transformer_model.pth')

# %% Train RL Model
# Set the appropriate hyperparameters
rl_model_hyperparams = {
    "epochs": 500,
    "batch_size": 32,
    "learning_rate": 0.001,
    "gamma": 0.99,  # discount factor
    "epsilon_start": 1.0,  # exploration rate at the beginning
    "epsilon_end": 0.01,  # minimum exploration rate
    "epsilon_decay": 0.995,  # exponential decay rate for exploration probability
}
train_rl(trading_agent, data, rl_model_hyperparams)

# %% Save RL Model
torch.save(rl_model.state_dict(), './models/rl_model.pth')