import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor

class Zenith(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)
        self.SetCash(10000)
        
        self.tickers = ["SPY", "QQQ", "SCHA", "VBK", "BOND", "BLV", "BSV", "GLD", "IAU"]
        self.stocks = []
        
        self.leverage = 4

        for i in self.tickers:
            security = self.AddEquity(i, Resolution.Hour)
            self.Securities[i].SetLeverage(self.leverage)
            self.Securities[i].SetDataNormalizationMode(DataNormalizationMode.Raw)
            security.SetDataNormalizationMode(DataNormalizationMode.Raw)
            self.stocks.append(security.Symbol)
            
            
        risk1 = self.GetParameter("risk-up")
        risk2 = self.GetParameter("risk-down")
        
        
        self.UniverseSettings.Resolution = Resolution.Hour
        # self.UniverseSettings.Leverage = self.leverage
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        #self.AddRiskManagement(TrailingStopRiskManagementModel(0.03))

            
        self.lookback = 60
        self.Trade = True

        self.risk_up = float(risk1)
        self.risk_down = float(risk2)
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday),
                 self.TimeRules.At(10, 0),
                 self.TradeSet)
                 
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Thursday),
                 self.TimeRules.At(10, 0),
                 self.Manage)
                 
    def Manage(self):
        
        for symbol in self.stocks:
            confidence = 1.0
            weight = 0.1
            
            #Takeprofit and Stop Loss
            if self.Portfolio[symbol].Invested and (self.Securities[symbol].Holdings.UnrealizedProfitPercent > self.risk_up or self.Securities[symbol].Holdings.UnrealizedProfitPercent < -(self.risk_down)):
                self.Liquidate(symbol)
        
    def OnData(self, data):
        

        if self.Trade == True:
            
            history = self.History(self.stocks, self.lookback, Resolution.Daily)
            self.historical_prices = {}
            self.predicted_prices = {}
            
            for symbol in self.stocks:
                if not history.empty and data.ContainsKey(symbol):
                    self.historical_prices[symbol] = list(history.loc[symbol.Value]['close'])
            
            for symbol in self.stocks:
                if symbol in self.historical_prices:
                    Y = []
                    X = []
                    for price in self.historical_prices[symbol]:
                        Y.append([price])
                        X.append([price])
                    
                    train_X, val_X, train_y, val_y = train_test_split(X, Y, random_state=1)
                    
                    price_model = RandomForestRegressor(random_state=1)
                    price_model.fit(train_X, train_y)
                    val_predictions = price_model.predict(val_X)
                    
                    self.predicted_prices[symbol] = val_predictions[4]
        else:
            return
        
        if not self.predicted_prices:
            return
        
        for symbol in self.stocks:
            if data.ContainsKey(symbol):
                
                magnitude = (self.predicted_prices[symbol] - self.Securities[symbol].Price) / self.Securities[symbol].Price
                confidence = 1.0
                weight = 0.1
                
                if magnitude > 0.001:
                    confidence = 0.75
                elif magnitude <= 0.001 and magnitude >= -0.001:
                    confidence = 0.5
                elif magnitude < -0.001:
                    confidence = 0.25
                
                if not self.predicted_prices[symbol] > (self.Securities[symbol].Price)*1.15:
                    self.SetHoldings(symbol, weight*self.leverage)
                elif self.predicted_prices[symbol] < self.Securities[symbol].Price:
                    self.Liquidate(symbol)
                elif self.predicted_prices[symbol]*1.1 < (self.Securities[symbol].Price):
                    self.SetHoldings(symbol, -weight*self.leverage)

        self.Trade = False

    def TradeSet(self):
        self.Trade = True