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

from System import *
from QuantConnect import *
from QuantConnect.Orders import *
from QuantConnect.Securities import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect import Resolution
from QuantConnect.Algorithm.Framework.Alphas import *
from AlgorithmImports import *
import torch
import torch.nn.functional as F

class NeuralNetworkAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)  # Set Start Date
        self.SetEndDate(2017, 12, 1) # Set End Date
        
        self.SetCash(1000000)  # Set Strategy Cash

        self.SetExecution(ImmediateExecutionModel())
        ## Initialize the various variables/helpers we'll need
        self.lastMonth = -1
        self.longs = []
        self.shorts = []
    

        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.SetRiskManagement(MaximumDrawdownPercentPortfolio(0.10))
        self.AddRiskManagement(MaximumDrawdownPercentPerSecurity(0.10))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
                

        # add symbol
        IAU = self.AddEquity("IAU", Resolution.Minute) #iShares Gold Trust
        #GLL = self.AddEquity("GLL", Resolution.Daily) #ProShares UltraShort Gold
        
        #USO = self.AddEquity("USO", Resolution.Daily) #United States Oil Fund LP
        #SCO = self.AddEquity("SCO", Resolution.Daily) #ProShares UltraShort Bloomberg Crude Oil
        
        #DBA = self.AddEquity("DBA", Resolution.Daily) #Invesco DB Agriculture Fund
        
        SPY = self.AddEquity("SPY", Resolution.Minute) #SPDR S&P 500 ETF Trust
        QQQ = self.AddEquity("QQQ", Resolution.Minute) #ProShares Short QQQ
        
        IEF = self.AddEquity("IEF", Resolution.Minute) #I-Shares 7-10 years treasury bond ETF
       #TLT = self.AddEquity("TLT", Resolution.Daily) #iShares 20 Plus Year Treasury Bond ETF
        #SPTL = self.AddEquity("SPTL", Resolution.Daily) #SPDR Portfolio Long Term Treasury ETF
       #CWB = self.AddEquity("CWB", Resolution.Daily) #SPDR Bloomberg Barclays Convertible Securities ETF
      # TMF = self.AddEquity("TMF", Resolution.Daily) #Direxion Daily 20+ Year Treasury Bull 3X Shares
      # TBT = self.AddEquity("TBT", Resolution.Daily) #ProShares UltraShort 20+ Year Treasury
      # MBB = self.AddEquity("MBB", Resolution.Daily) #iShares MBS ETF
        
     #  ANGL = self.AddEquity("ANGL", Resolution.Daily) #VanEck Vectors Fallen Angel High Yield Bond ETF
     #  SJNK = self.AddEquity("SJNK", Resolution.Daily) #SPDR Bloomberg Barclays Short Term High Yield Bond ETF
     #  HYGV = self.AddEquity("HYGV", Resolution.Daily) #FlexShares High Yield Value-Scored US Bond Index Fund
        
        
      # VNQ = self.AddEquity("VNQ", Resolution.Daily) #Vanguard Real Estate ETF
     #  IYR = self.AddEquity("IYR", Resolution.Daily) #iShares U.S. Real Estate ETF
        ITB = self.AddEquity("ITB", Resolution.Minute) #iShares U.S. Home Construction ETF
        
      # IVV = self.AddEquity("IVV", Resolution.Daily) #iShares Core S&P 500 ETF
      # VTI = self.AddEquity("VTI", Resolution.Daily) #Vanguard Total Stock Market ETF
     #  VGT = self.AddEquity("VGT", Resolution.Daily) #Vanguard Information Technology ETF
     #  XLK = self.AddEquity("XLK", Resolution.Daily) #Technology Select Sector SPDR Fund
      # SOXX = self.AddEquity("SOXX", Resolution.Daily) #iShares Semiconductor ETF
        GXC = self.AddEquity("GXC", Resolution.Minute) #SPDR S&P China ETF
        
     #  IHI = self.AddEquity("IHI", Resolution.Daily) #iShares U.S. Medical Devices ETF
        
        self.reference = "SPY"
        self.SetBenchmark("SPY")  
        self.SetWarmUp(50) 
                
        self.symbols = [IAU.Symbol, SPY.Symbol, QQQ.Symbol, IEF.Symbol, ITB.Symbol, GXC.Symbol]  #GLL.SymboSPY.Symbol, USO.Symbol, SCO.Symbol, DB QQQ.Symbol, IEF.Symbol, 
                     #   TLT.Symbol, SPTL.Symbol,CWB.Symbol,TMF.Symbol,TBT.Symbol,MBB.Symbol, ANGL.Symbol, SJNK.Symbol,
                    #    HYGV.Symbol, VNQ.Symbol, IYR.Symbol, ITB.Symbol, IVV.Symbol, VTI.Symbol,
                   #     VGT.Symbol, XLK.Symbol,SOXX.Symbol, GXC.Symbol, IHI.Symbol 
                     #  ] #SVXY.Symbol,  # using a list can extend to condition for multiple symbols
       
        
        self.lookback = 30 # days of historical data (look back)
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 28), self.NetTrain) # train the NN
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 30), self.Trade)
    
    
    def OnData(self, data):
        pass

    
    def NetTrain(self):
        # Daily historical data is used to train the machine learning model
        history = self.History(self.symbols, self.lookback + 1, Resolution.Minute)
        
        # 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=10, n_output=1)     # define the network
                optimizer = torch.optim.SGD(net.parameters(), lr=0.2)
                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 
        ''' 
        for holding in self.Portfolio.Values:
            if self.CurrentSlice[holding.Symbol].Open < self.sell_prices[holding.Symbol] and holding.Invested:
                self.Liquidate(holding.Symbol)
            
            if self.CurrentSlice[holding.Symbol].Open > self.buy_prices[holding.Symbol] and not holding.Invested:
                self.SetHoldings(holding.Symbol, 1 / len(self.symbols))

            if self.CurrentSlice[holding.Symbol].Open is self.buy_prices[holding.Symbol] and not holding.Invested:
                self.SetHoldings(holding.Symbol, 1 / len(self.symbols))

            
# 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