from arch import arch_model
import numpy as np

class StrategyManagerUtility:
    def __init__(self, quantScenario):
        self.quantScenario = quantScenario
    
    def setupModel(self, preference):
        self.preference = 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_gm = arch_model(preference.garchData, p = 1, q = 1, o = 1, vol = 'GARCH', dist = 't', rescale = "FALSE")
            
            egarch_result = egarch_gm.fit()
                
            # 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 isnan(self, value):
        try:
            import math
            return math.isnan(float(value))
        except:
            return False
    
    def genericGarchLog(self, preference, inputGm_forecast, garchOutput):
        
        sdVariance = (preference.egarchMean)/garchOutput[0]
        adjustedSdVariance = (preference.egarchMean)/garchOutput[0]
        
        #formatting the data for logging below: removing Nan values, rounding very small numbers up to zero, and very large numbers down to 100%
        if self.isnan(sdVariance):
            return
        
        if "e+" in str(sdVariance):
            sdVariance = 0
        
        if sdVariance > 1:
                adjustedSdVariance = 100

        sdVarianceNumpyConversion = round(float(sdVariance * 100), 4)
        
        preference.quantScenario.Debug(str(preference.quantScenario.Time) + " - " + str(preference.garchData[-1]))
        preference.quantScenario.Debug("SOXL Allocation(raw): " + str(sdVariance) + ", Adjusted Allocation: "+ str(adjustedSdVariance) + "%, Percentage Conversion: " + str(sdVarianceNumpyConversion) + "%")
        #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("Element in Garch Data Array: "+ str(preference.garchData[:5]))
        preference.quantScenario.Debug("length ---> " + str(len(self.preference.garchData)))

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-04-01 13:" in str(self.Time):
                self.investorPrefernce.injectTodaysStandardDeviationValue(3.5)
        if "2021-04-01 14:" in str(self.Time):
                self.investorPrefernce.injectTodaysStandardDeviationValue(5.34)

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, 10, 20) #stock split is 2015,5,20 - 50%
        self.quantScenario.SetEndDate(2021, 12, 31)
        self.quantScenario.SetCash(20000)  # Set Strategy Cash
        self.leveragedETFSymbol = "TQQQ"
        self.tradedETFSymbol = "SOXL"
        self.minutesBeforeMarketClose = 1
        self.counter = 0
        self.investmentInterval = 14
        self.esppInterval = 90
        self.indicatorPeriod = 7
        self.garchData = []
        self.egarchMean = 0
        list = []
        self.egarchMeanArray = np.array(list)
        self.resolution = Resolution.Daily
        self.stockSplitMapping = {29.72511: 4.75, 
        38.55771: 5.36, 
        41.87474: 5.24,
        42.2561: 3.83,
        41.22449: 2.99,
        34.39242: 4.47 }
        
    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.Adjusted)
        
        self.tradedETF = self.quantScenario.AddEquity(self.tradedETFSymbol, self.resolution)
        self.tradedETF.SetDataNormalizationMode(DataNormalizationMode.Adjusted)
        self.tradedETF.SetLeverage(1)
        
        #without this code, injecting into the algorithm will not work.
        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 == 0:
            investment = 1200
            #1230 - indiviudal
            #9859 - group - 3000 Justin, 1230
            #Rigby's yearly contribution: 115000
            self.quantScenario.Portfolio.SetCash(self.quantScenario.Portfolio.Cash + investment)
        
        if self.counter % self.esppInterval == 0:
            investment = 0
            #4500
            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 -------------------------------------------------------------
        roundedStdIndicatorValue = round(self.quantScenario.stdIndicator.Current.Value, 5)
        if roundedStdIndicatorValue in self.stockSplitMapping.keys():
            adjustedStockSplitValue = self.stockSplitMapping[roundedStdIndicatorValue]
            self.quantScenario.stdIndicator.Current.Value = adjustedStockSplitValue
            self.garchData.append(adjustedStockSplitValue)
        else:
            self.garchData.append(roundedStdIndicatorValue)
        
    def dailyComputation(self):
        
        self.normalizeGarchDataFromStockSplit() #--------------------------------remove this line with a conditional
            
        self.counter += 1
        
        self.rebalanceGarchArray()
            
        self.strategyManager.setupModel(self)
        
        #self.Debug("Garch Data: " + str(data))
    
    def rebalanceGarchArray(self):
        if len(self.garchData) > 60:
            self.garchData.pop(0)
    
    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))