from arch import arch_model
import numpy as np

class StrategyManagerUtility:
    def __init__(self, quantScenario):
        self.quantScenario = quantScenario
    
    def setupModel(self, preference):
    # GJR GARCH mode

        # Attempt to Fit the model
        try:    
            egarch_gm = arch_model(preference.garchData, p = 1, q = 1, o = 1, vol = 'GARCH', dist = 't')
                
            egarch_result = egarch_gm.fit(disp = 'off')
                
            # Make 1-period ahead forecast
            gm_forecast = egarch_result.forecast(horizon = 1)
            preference.egarchMean = gm_forecast.mean.values[-1 ]
            if preference.egarchMean:
                preference.egarchMeanArray = np.append( preference.egarchMeanArray, preference.egarchMean)
            else:
                pass
            
            garchOutput = gm_forecast.variance[-1:].values
            
            #printing the logs
            self.genericGarchLog(preference, gm_forecast, garchOutput[0])
            
            if preference.quantScenario.IsWarmingUp:
                return
            else:
                self.setHoldingsBasedOnEGarch(garchOutput[0], preference)
        except Exception:
            return
        
        #self.Debug("----------> " + str(self.gm_result.summary()))
    
    def setHoldingsBasedOnEGarch(self, garchOutput, preference):
        
        self.multipler = 2
        
        if garchOutput[0] < preference.egarchMean:
            preference.quantScenario.SetHoldings(preference.tradedETFSymbol, 1)
        elif garchOutput[0] > preference.egarchMean*self.multipler and  garchOutput[0] > 0:
            #self.SetHoldings(self.tradedETFSymbol, (self.egarchMean)/garchOutput[0])
            sdVariance = (preference.egarchMean)/garchOutput[0]
            
            if sdVariance > 1:
                sdVariance = 1
                
            preference.quantScenario.SetHoldings(preference.tradedETFSymbol, sdVariance)
            #self.SetHoldings([PortfolioTarget("TQQQ", 0), PortfolioTarget("ETF", int(sdVariance))], True)
            
    def genericGarchLog(self, preference, inputGm_forecast, garchOutput):
        
        sdVariance = (preference.egarchMean)/garchOutput[0]
        
        preference.quantScenario.Debug(str(preference.quantScenario.Time) + " - " + str(preference.garchData[-1]))
        preference.quantScenario.Debug("SOXL Allocation: "+ str(sdVariance))
        preference.quantScenario.Debug("Mean, based on Garch Data: "+ str(preference.egarchMean))
        preference.quantScenario.Debug("Forecast Variance, based on Garch Data: "+ str(inputGm_forecast.variance[-1:].values))
        preference.quantScenario.Debug("Yesterday's Standard Deviation: "+str(preference.garchData[-1]))
        preference.quantScenario.Debug("")
        #preference.quantScenario.Debug("Element in Garch Data Array: "+ str(preference.garchData[:5]))

from InvestorPreferenceUtility import InvestorPreferenceUtility

class SOXLStandardDeviation(QCAlgorithm):
    
    #101293240057
    
    def Initialize(self):
        
        self.investorPrefernce = InvestorPreferenceUtility(self)    
    
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        if "2021-01-22 15:" in str(self.Time):
            self.investorPrefernce.injectTodaysStandardDeviationValue(5)

from StrategyManagerUtility import StrategyManagerUtility
import numpy as np
    
class InvestorPreferenceUtility:
    
    def __init__(self, quantconnectObject):
        self.quantScenario = quantconnectObject
        self.strategyManager = StrategyManagerUtility(self)
        
        self.initalizeAttributes()
        self.setupIndicators()
        self.setupScheduler()
        
    def initalizeAttributes(self):
        self.quantScenario.SetStartDate(2020, 12, 1)
        self.quantScenario.SetEndDate(2021, 1, 25) 
        self.quantScenario.SetCash(2000)  # Set Strategy Cash
        self.leveragedETFSymbol = "TQQQ"
        self.tradedETFSymbol = "SOXL"
        self.minutesBeforeMarketClose = 1
        self.counter = 0
        self.investmentInterval = 14
        self.indicatorPeriod = 7
        self.garchData = []
        self.egarchMean = 0
        list = []
        self.egarchMeanArray = np.array(list)
        self.resolution = Resolution.Daily
        
    def setupIndicators(self):
        #Warming up the function
        self.warmUpTimeForBacktestScenario = self.indicatorPeriod
        self.quantScenario.SetWarmUp(timedelta(days=self.indicatorPeriod))
        
        self.leveragedETF = self.quantScenario.AddEquity(self.leveragedETFSymbol, self.resolution)
        self.leveragedETF.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.tradedETF = self.quantScenario.AddEquity(self.tradedETFSymbol, self.resolution)
        self.tradedETF.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.forceQuantToCallOnDataEachHour = self.quantScenario.AddEquity("VOO", Resolution.Hour)
        
        #Setup standard deviation indicator
        self.quantScenario.stdIndicator = self.quantScenario.STD(self.leveragedETFSymbol, self.indicatorPeriod, Resolution.Daily)
    
    def setupScheduler(self):
        self.quantScenario.Schedule.On(self.quantScenario.DateRules.EveryDay(self.leveragedETFSymbol), self.quantScenario.TimeRules.AfterMarketOpen(self.leveragedETFSymbol, 1), self.dailyComputation)
        
        self.quantScenario.Schedule.On(self.quantScenario.DateRules.EveryDay(self.leveragedETFSymbol), self.quantScenario.TimeRules.BeforeMarketClose(self.leveragedETFSymbol, self.minutesBeforeMarketClose), self.dollarCostAverage)

    def dollarCostAverage(self):
        self.counter += 1
        
        if self.counter % self.investmentInterval:
            investment = 1000
            self.quantScenario.Portfolio.SetCash(self.quantScenario.Portfolio.Cash + investment)
    
    def normalizeGarchDataFromStockSplit(self):
        #This is a temporary conditional statement to account for the price split on 1/22/2021, please remove this conditional statement after 60 days from 1/22/2021 -------------------------------------------------------------
        if self.quantScenario.stdIndicator.Current.Value == 29.7251056536802:
            self.garchData.append(4.75123456789)
        else:
            self.garchData.append(self.quantScenario.stdIndicator.Current.Value)
        
    
    def dailyComputation(self):
        
        self.normalizeGarchDataFromStockSplit() #--------------------------------remove this line with a conditional
            
        self.counter += 1
        
        self.strategyManager.setupModel(self)
        
        if len(self.garchData) > 60:
            self.garchData.pop(0)
        
        #self.Debug("Garch Data: " + str(data))
    
    def injectTodaysStandardDeviationValue(self, inputStandardDeviation):
        
        self.quantScenario.Debug("----------->Injecting Stardard Deviation Data for the next day")
        self.garchData.append(inputStandardDeviation)
        self.strategyManager.setupModel(self)
        #mean =  np.mean(preference.egarchMeanArray)
        #preference.quantScenario.Debug("egarch mean: " + str(mean))