# region imports
from AlgorithmImports import *
import torch
from torch import nn
import joblib
# endregion

class PyTorchExampleAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2022, 7, 4)
        self.SetCash(100000)
        self.symbol = self.AddEquity("SPY", Resolution.Daily).Symbol

        training_length = 252*2
        self.training_data = RollingWindow[float](training_length)
        history = self.History[TradeBar](self.symbol, training_length, Resolution.Daily)
        for trade_bar in history:
            self.training_data.Add(trade_bar.Close)

        if self.ObjectStore.ContainsKey("model"):
            file_name = self.ObjectStore.GetFilePath("model")
            self.model = joblib.load(file_name)
        else:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            self.model = NeuralNetwork().to(device)
            self.Train(self.my_training_method)
            
        self.Train(self.DateRules.Every(DayOfWeek.Sunday), self.TimeRules.At(8,0), self.my_training_method)
        
    def get_features_and_labels(self, n_steps=5):
        close_prices = list(self.training_data)[::-1]

        features = []
        labels = []
        for i in range(len(close_prices)-n_steps):
            features.append(close_prices[i:i+n_steps])
            labels.append(close_prices[i+n_steps])
        features = np.array(features)
        labels = np.array(labels)

        return features, labels

    def my_training_method(self):
        features, labels = self.get_features_and_labels()

        # Set the loss and optimization functions
        # In this example, use the mean squared error as the loss function and stochastic gradient descent as the optimizer
        loss_fn = nn.MSELoss()
        learning_rate = 0.001
        optimizer = torch.optim.SGD(self.model.parameters(), lr=learning_rate)
        
        # Create a for-loop to train for preset number of epoch
        epochs = 5
        for t in range(epochs):
            # Create a for-loop to fit the model per batch
            for batch, (feature, label) in enumerate(zip(features, labels)):
                # Compute prediction and loss
                pred = self.model(feature)
                real = torch.from_numpy(np.array(label).flatten()).float()
                loss = loss_fn(pred, real)
            
                # Perform backpropagation
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

    def OnData(self, slice: Slice) -> None:
        if self.symbol in slice.Bars:
            self.training_data.Add(slice.Bars[self.symbol].Close)

        features, __ = self.get_features_and_labels()
        prediction = self.model(features[-1].reshape(1, -1))
        prediction = float(prediction.detach().numpy()[-1])

        if prediction > slice[self.symbol].Price:
            self.SetHoldings(self.symbol, 1)
        elif prediction < slice[self.symbol].Price:            
            self.SetHoldings(self.symbol, -1)

    def OnEndOfAlgorithm(self):
        model_key = "model"
        file_name = self.ObjectStore.GetFilePath(model_key)
        joblib.dump(self.model, file_name)
        self.ObjectStore.Save(model_key)

class NeuralNetwork(nn.Module):
    # Model Structure
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(5, 5),   # input size, output size of the layer
            nn.ReLU(),         # Relu non-linear transformation
            nn.Linear(5, 5),
            nn.ReLU(),  
            nn.Linear(5, 1),   # Output size = 1 for regression
        )
    
    # Feed-forward training/prediction
    def forward(self, x):
        x = torch.from_numpy(x).float()   # Convert to tensor in type float
        result = self.linear_relu_stack(x)
        return result