from arch import arch_model
import numpy as np

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

        # 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
            
            #printing the logs
            self.genericGarchLog(preference, gm_forecast)
                
            garchOutput = gm_forecast.variance[-1:].values
            
            if preference.quantScenario.IsWarmingUp:
                return
            else:
                self.setHoldingsBasedOnEGarch(garchOutput[0], preference)
        except Exception:
            pass
        
        #self.Debug("----------> " + str(self.gm_result.summary()))
    
    def setHoldingsBasedOnEGarch(self, garchOutput, preference):
        
        self.multipler = 1
        
        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]
            preference.quantScenario.SetHoldings(preference.tradedETFSymbol, sdVariance)
            #self.SetHoldings([PortfolioTarget("TQQQ", 0), PortfolioTarget("ETF", int(sdVariance))], True)
            
    def genericGarchLog(self, preference, inputGm_forecast):
        
        preference.quantScenario.Debug(str(preference.quantScenario.Time))
        preference.quantScenario.Debug("Mean, based on variance: "+ str(preference.egarchMean))
        preference.quantScenario.Debug("Forecast Variance based on Garch Data: "+ str(inputGm_forecast.variance[-1:].values))
        preference.quantScenario.Debug("Standard Deviation, today's Garch element: "+str(preference.garchData[-1]))
        preference.quantScenario.Debug("Garch Data: "+ str(preference.garchData[:5]))

from InvestorPreferenceUtility import InvestorPreferenceUtility

class SOXLStandardDeviation(QCAlgorithm):
    
    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
        '''

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, 1, 1)
        self.quantScenario.SetEndDate(2020, 12, 31) 
        self.quantScenario.SetCash(100000)  # Set Strategy Cash
        self.leveragedETFSymbol = "VOO"
        self.tradedETFSymbol = "UPRO"
        self.minutesBeforeMarketClose = 1
        self.counter = 0
        self.investmentInterval = 14
        self.indicatorPeriod = 7
        self.garchData = []
        self.priceHistory = []
        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(self.warmUpTimeForBacktestScenario)
        
        #1. Update the AddEquity command to request TQQQ data
        self.leveragedETF = self.quantScenario.AddEquity(self.leveragedETFSymbol, self.resolution)
        self.leveragedETF.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        #1. Update the AddEquity command to request SOXL data
        self.tradedETF = self.quantScenario.AddEquity(self.tradedETFSymbol, self.resolution)
        self.tradedETF.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.hedge = self.quantScenario.AddEquity("SQQQ", self.resolution)
        self.hedge.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        #Setup standard deviation indicator
        self.quantScenario.stdIndicator = self.quantScenario.STD(self.leveragedETFSymbol, self.indicatorPeriod, self.resolution)
    
    def setupScheduler(self):
        self.quantScenario.Schedule.On(self.quantScenario.DateRules.EveryDay(self.leveragedETFSymbol), self.quantScenario.TimeRules.BeforeMarketClose(self.leveragedETFSymbol, self.minutesBeforeMarketClose), 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 = 1666
            self.quantScenario.Portfolio.SetCash(self.quantScenario.Portfolio.Cash + investment);
    
    def dailyComputation(self):
        
        #calculate the percentage change
        self.garchData.append(self.quantScenario.stdIndicator.Current.Value)
        
        self.counter += 1
        
        self.strategyManager.setupModel(self)
        
        #if "2020-02-10" in str(self.quantScenario.Time):
        #    self.injectTodaysStandardDeviationValue(1.2)
        
        if len(self.garchData) > 60:
            self.garchData.pop(0)
        
        #self.Debug("Garch Data: " + str(data))
    
    def injectTodaysStandardDeviationValue(self, inputStandardDeviation):
        
        self.garchData.append(inputStandardDeviation)
        self.strategyManager.setupModel(self)
        #mean =  np.mean(preference.egarchMeanArray)
        #preference.quantScenario.Debug("egarch mean: " + str(mean))