from AlgorithmImports import *


class TransdimensionalPrism(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        #self.SetEndDate(2020,1,1)
        self.SetCash(12000)  # Set Strategy Cash
        self.SetWarmUp(1000)
        
        #SPY DATA
        spy = self.AddEquity("SPY", Resolution.Minute)
        spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.spy = spy.Symbol
        self.spyCallcontract = None
        self.spyPutcontract = None
        self.spySmallPutcontract = None
        
        #SOXL DATA
        soxl = self.AddEquity("SOXL", Resolution.Minute)
        soxl.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.soxl = soxl.Symbol
        self.soxlCallcontract = None        
        
        #QQQ
        qqq = self.AddEquity("QQQ", Resolution.Minute)
        qqq.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.qqq = qqq.Symbol
        self.qqqCallcontract = None
        self.qqqPutcontract = None
        self.qqqSmallPutcontract = None
        
        #TQQQ Data
        #LEVERAGED ETFS: [SPXL:Leveraged SPY, TQQQLeveraged TQQQ, SOXL]
        tqqq = self.AddEquity("TQQQ", Resolution.Minute)
        tqqq.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.tqqq = tqqq.Symbol
        self.tqqqCallcontract = None
        self.tqqqPutcontract = None
        self.tqqqSmallPutcontract = None
        
        #BND
        bnd = self.AddEquity("BND", Resolution.Minute)
        self.bnd = bnd.Symbol
        
        #QQQ
        qqq = self.AddEquity("QQQ", Resolution.Minute)
        self.qqq = qqq.Symbol

        
        #Scheduled Events============================================
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 1), self.monthlyRebalance)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 1), self.captureOpening)
        #Variables needed for stoploss
        self.stoplosshold = 0
        
        self.dailyOpen: Dict[Symbol,float] = {
            self.spy: 0,
        }

        self.weights: Dict[Symbol,float] = {
            self.bnd:.8,
        }
        
        #Underlyings and Options dictionary
        #self.uNos = { self.spy  : self.spycontract, #self.tqqq : self.tqqqcontract }
        self.uNos = {
                     'call'     : {self.spy  : self.spyCallcontract, 
                                   self.qqq : self.qqqCallcontract
                     },
                     'put'      : {self.spy   : self.spyPutcontract},
                     'smallPut' : {self.qqq  : self.qqqSmallPutcontract}
                    }
        

    def OnData(self, data): 
        
        if not self.Portfolio[self.bnd].Invested:
            self.SetHoldings(self.bnd, 0.8)
            
        # DO HEDGE
        self.getContracts()
        
        self.timeDC()
        
        self.setContracts()
        
        self.exerciseContracts()
        
        
        
    def getContracts(self):
        '''
        \nThis method is used to loop through the nested 
        \n dictionary: self.uNos (underlyings and options)

        '''
        if self.IsWarmingUp:
            return
        for right, cluster in self.uNos.items():
            if right == 'call':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        # ATM call option 
                        targetStrike = (self.Securities[underlying].Price *1.2) - (self.Securities[underlying].Price * 1.2)%5
                        xcontracts = self.OptionChainProvider.GetOptionContractList(underlying, self.Time)
                        calls = [x for x in xcontracts if x.ID.OptionRight == OptionRight.Call]
                        calls = sorted( sorted(calls, key = lambda x: x.ID.Date, reverse = True),
                        key = lambda x: x.ID.StrikePrice)
                        calls = [x for x in calls if x.ID.StrikePrice == targetStrike]
                        calls = [x for x in calls if 250 < (x.ID.Date - self.Time).days]
                        if len(calls) == 0:
                            calls = sorted(calls, key = lambda x: x.ID.Date, reverse=True)
                        if len(calls) == 0:    
                            return None
                        self.AddOptionContract(calls[0], Resolution.Minute)
                        self.uNos[right][underlying] = calls[0]
                        
                        
            if right == 'put':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        #40% OTM call option 
                        targetStrike = (self.Securities[underlying].Price * .6) - (self.Securities[underlying].Price * .6)%5
                        xcontracts = self.OptionChainProvider.GetOptionContractList(underlying, self.Time)
                        puts = [x for x in xcontracts if x.ID.OptionRight == OptionRight.Put]
                        puts = sorted( sorted(puts, key = lambda x: x.ID.Date, reverse = True),
                        key = lambda x: x.ID.StrikePrice)
                        puts = [x for x in puts if x.ID.StrikePrice == targetStrike]
                        puts = [x for x in puts if 300 < (x.ID.Date - self.Time).days <= 420]
                        if len(puts) == 0:
                            return None
                        self.AddOptionContract(puts[0], Resolution.Minute)
                        self.uNos[right][underlying] = puts[0]
                        
            if right == 'smallPut':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        #40% OTM call option 
                        targetStrike = (self.Securities[underlying].Price * .8) - (self.Securities[underlying].Price * .8)%5
                        xcontracts = self.OptionChainProvider.GetOptionContractList(underlying, self.Time)
                        puts = [x for x in xcontracts if x.ID.OptionRight == OptionRight.Put]
                        puts = sorted( sorted(puts, key = lambda x: x.ID.Date, reverse = True),
                        key = lambda x: x.ID.StrikePrice)
                        puts = [x for x in puts if x.ID.StrikePrice == targetStrike]
                        puts = [x for x in puts if 300 < (x.ID.Date - self.Time).days]
                        if len(puts) == 0:
                            return None
                        self.AddOptionContract(puts[0], Resolution.Minute)
                        self.uNos[right][underlying] = puts[0]
            
                 

    def timeDC(self):
        
        for right, cluster in self.uNos.items():
            if right == 'call':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    if (contract.ID.Date - self.Time).days < 20:
            #self.Log("Liquidated time 180")
                        self.Liquidate(contract)
                        self.RemoveSecurity(contract)
                        self.uNos[right][underlying] = None
            if right == 'put':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    if (contract.ID.Date - self.Time).days < 30:
            #self.Log("Liquidated time 180")
                        self.Liquidate(contract)
                        self.RemoveSecurity(contract)
                        self.uNos[right][underlying] = None
                return
    
    
    def setContracts(self):
        
        for right, cluster in self.uNos.items():
            if right == 'call':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    if not self.Portfolio[contract].Invested:
                        if self.CurrentSlice.Bars.ContainsKey(underlying):
                            #self.SetHoldings(contract, 0.1)
                            self.MarketOrder(contract, 1)
                            
            if right == 'put':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    if not self.Portfolio[contract].Invested:
                        #self.SetHoldings(contract, 0.03)
                        if self.Securities[contract].Price != 0:
                            # SHILE: PUT ADDITIONAL CHECKS HERE
                            self.Log(self.Securities[contract].Price)
                            self.MarketOrder(contract, 1)
                        
            if right == 'smallPut':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    if not self.Portfolio[contract].Invested:
                        #self.SetHoldings(contract, 0.07)
                        if self.Securities[contract].Price != 0:
                            # SHILE: PUT ADDITIONAL CHECKS HERE
                            self.MarketOrder(contract, 1)

    
    def exerciseContracts(self):
        
        for right, cluster in self.uNos.items():
            if right == 'call':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    
                    #Log the unrealized potential of the call option
                    #If the unrealized profit of the option contract has increased by 100%
                    if self.Portfolio[contract].UnrealizedProfitPercent > 1.5:
                        #self.Log(self.Securities[underlying].Price)
                        self.Liquidate(contract)
                        #self.RemoveSecurity(contract)
                        self.RemoveOptionContract(contract)
                        self.uNos[right][underlying] = None
            if right == 'put':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    # For the large scale puts, if the unrealized profit has increased by 310%
                    # This can be considered the 'F#%#@ Me' moment
                    if self.Portfolio[contract].UnrealizedProfitPercent > 2.5:    
                        self.Liquidate(contract)
                        #self.RemoveSecurity(contract)
                        self.RemoveOptionContract(contract)
                        self.uNos[right][underlying] = None
                
            if right == 'smallPut':
                for underlying, contract in self.uNos[right].items():
                    if contract is None:
                        continue
                    # IF the unrealized profit of the smaller scale put has increased by 100%
                    if self.Portfolio[contract].UnrealizedProfitPercent > 1:    
                        self.Liquidate(contract)
                        #self.RemoveSecurity(contract)
                        self.RemoveOptionContract(contract)
                        self.uNos[right][underlying] = None
                
                
                
    def monthlyRebalance(self):
        # Rebalance portfolio monthly 
        if self.IsWarmingUp:
            return
        for key,value in self.weights.items():
            self.SetHoldings(key, value)
            
        
    def captureOpening(self):
        #Grabs the daily opening price of spy for our stoploss method
        if self.IsWarmingUp:
            return
        for key, value in self.dailyOpen.items():
            if self.CurrentSlice.Bars.ContainsKey(key):
                self.dailyOpen[key] = self.CurrentSlice[key].Open
            self.stoplosshold = 0
            
            
    def stoploss(self, data):
        '''
        Stoploss logic:
            - If spy drops more than 5% liquidate entire equity portfolio
            - Change stoplosshold value to 1, this indicates that the portfolios SL has been hit 
            and were going to hold until the next trading day
        '''
        if self.IsWarmingUp:
            return
        
        
        for symbol, weight in self.weights.items():
            if self.Securities[symbol].Price == 0:
                continue
            open = self.dailyOpen[symbol]
            #See if any symbol has decreased more than 5% in a given day, liquidate if true and check next one... 
            if ((self.Securities[symbol].Price-open)/self.Securities[symbol].Price) < -.05:
                self.SetHoldings(symbol, 0)
                
                self.stoplosshold = 1
                #self.Log('HIT')