import clr
clr.AddReference("System")
clr.AddReference("QuantConnect.Algorithm")
clr.AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *

import numpy as np
import torch
import torch.nn.functional as F

class PytorchNeuralNetworkAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 5, 2)  # Set Start Date
        #self.SetEndDate(2020, 10, 8) # Set End Date
        self.SetBrokerageModel(BrokerageName.OandaBrokerage)
        self.SetCash(1000000)  # Set Strategy Cash
        
        # add symbol
        spy = self.AddForex("EURUSD", Resolution.Minute, Market.Oanda) #eurchf eurnzd eurjpy
        spy1 = self.AddForex("USDCHF", Resolution.Minute, Market.Oanda)
        gug = self.AddForex("GBPUSD", Resolution.Minute, Market.Oanda)
        spy2 = self.AddForex("EURGBP", Resolution.Minute, Market.Oanda)
        spy3 = self.AddForex("AUDUSD", Resolution.Minute, Market.Oanda)
        spy4 = self.AddForex("EURAUD", Resolution.Minute, Market.Oanda)
        spy5 = self.AddForex("AUDCHF", Resolution.Minute, Market.Oanda)
        spy6 = self.AddForex("USDJPY", Resolution.Minute, Market.Oanda)
        spy7 = self.AddForex("GBPJPY", Resolution.Minute, Market.Oanda)
        spy8 = self.AddForex("AUDJPY", Resolution.Minute, Market.Oanda)
        spy9 = self.AddForex("GBPAUD", Resolution.Minute, Market.Oanda)
        spy10 = self.AddForex("NZDUSD", Resolution.Minute, Market.Oanda)
        spy11 = self.AddForex("AUDNZD", Resolution.Minute, Market.Oanda)
        spy12 = self.AddForex("USDCAD", Resolution.Minute, Market.Oanda)
        spy13 = self.AddForex("EURNZD", Resolution.Minute, Market.Oanda)
        self.symbols = [spy.Symbol, spy1.Symbol, gug.Symbol, spy2.Symbol, spy3.Symbol, spy4.Symbol, spy5.Symbol,
                        spy6.Symbol, spy7.Symbol, spy8.Symbol, spy9.Symbol, spy10.Symbol, spy11.Symbol, spy12.Symbol, spy13.Symbol]# using a list can extend to condition for multiple symbols# , 
            
        self.lookback = 5 # days of historical data (look back)
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(timedelta(minutes=15)), self.NetTrain) #AfterMarketOpen("EURUSD", 180), self.NetTrain) # train the NN, TimeRules.AfterMarketOpen("EURUSD", 90,95),
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(timedelta(minutes=20)), self.Trade) #AfterMarketOpen("EURUSD", 200), self.Trade)
        slowperiod = 140
        self.SetWarmUp(slowperiod)
    def NetTrain(self):
        # Daily historical data is used to train the machine learning model
        history = self.History(self.symbols, self.lookback + 1, Resolution.Daily)
        
        # dicts that store prices for training
        self.prices_x = {} 
        self.prices_y = {}
        
        # dicts that store prices for sell and buy
        self.sell_prices = {}
        self.buy_prices = {}
        
        for symbol in self.symbols:
            if not history.empty:
                # x: preditors; y: response
                self.prices_x[symbol] = list(history.loc[symbol.Value]['open'])[:-1]
                self.prices_y[symbol] = list(history.loc[symbol.Value]['open'])[1:]
        
        for symbol in self.symbols:
            # if this symbol has historical data
            if symbol in self.prices_x:
                
                net = Net(n_feature=1, n_hidden=80, n_output=1)     # define the network
                optimizer = torch.optim.SGD(net.parameters(), lr=0.4)
                loss_func = torch.nn.MSELoss()  # this is for regression mean squared loss
                
                for t in range(200):
                    # Get data and do preprocessing
                    x = torch.from_numpy(np.array(self.prices_x[symbol])).float()
                    y = torch.from_numpy(np.array(self.prices_y[symbol])).float()
                    
                    # unsqueeze data (see pytorch doc for details)
                    x = x.unsqueeze(1) 
                    y = y.unsqueeze(1)
                
                    prediction = net(x)     # input x and predict based on x

                    loss = loss_func(prediction, y)     # must be (1. nn output, 2. target)

                    optimizer.zero_grad()   # clear gradients for next train
                    loss.backward()         # backpropagation, compute gradients
                    optimizer.step()        # apply gradients
            
            # Follow the trend    
            self.buy_prices[symbol] = net(y)[-1] + np.std(y.data.numpy())
            self.sell_prices[symbol] = net(y)[-1] - np.std(y.data.numpy())
        
    def Trade(self):
        ''' 
        Enter or exit positions based on relationship of the open price of the current bar and the prices defined by the machine learning model.
        Liquidate if the open price is below the sell price and buy if the open price is above the buy price 
        ''' 
        #if (not data.ContainsKey(self.symbols)): return
        for holding in self.Portfolio.Values:
            if self.CurrentSlice[holding.Symbol].Open < self.sell_prices[holding.Symbol] and not holding.Invested:
                self.SetHoldings(holding.Symbol, 5 / len(self.symbols))
                
            elif self.CurrentSlice[holding.Symbol].Open > self.buy_prices[holding.Symbol] and not holding.Invested:
                self.SetHoldings(holding.Symbol, -5 / len(self.symbols))
                
            elif holding.Invested:
                if self.CurrentSlice[holding.Symbol].Open < self.sell_prices[holding.Symbol] and holding.IsShort:
                    self.Liquidate(holding.Symbol)
                if self.CurrentSlice[holding.Symbol].Open > self.buy_prices[holding.Symbol] and holding.IsLong:
                    self.Liquidate(holding.Symbol)
                
        
# class for Pytorch NN model
class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden)   # hidden layer
        self.predict = torch.nn.Linear(n_hidden, n_output)   # output layer
    
    def forward(self, x):
        x = F.relu(self.hidden(x))      # activation function for hidden layer
        x = self.predict(x)             # linear output
        return x