'''
MULTI STOCK MODE

UNIVERSE -  Self Listed Stocks


DYNAMIC 

FROM the ---- > MANUAL / US / SPECIFIC COUNTRY EQUITIES 

SCAN FOR ---->  NEAR EARNINGS / INCREASED VOLATALITY / INCREASED VOLUME / INCREASED MOMENTUM OR TREND


STRATEGY SPECIFIC BACKTESTED UNIVERSE

FROM --> MANUAL / US / SPECIFIC COUNTRY EQUITIES 

SELECTED THOSE WHICH HAVE PERFORMED WELL IN LAST 6MONTHS TIME ZONE


disableHAcandles & use plan candles
disableCalls
disablePuts
disableShortStocks
allow or disablePostMarketHours
trailingstoplesspercentage-Options
trailingstoplesspercentage-Stocks
profittakingpercentage-Stocks
profittakingpercentage-Options
comparing returns with SPY in charts
options - itm - otm - atm 

IGNORE ERRORS OF ANY SPECIFIC STOCK FROM UNIVERSE - HAS NO DATA & CONTINUE PROGRAM

WEIGHTAGE ----> BUY SELL PERCENTAGE WEIGHTAGE - NOT NUMBER OF SHARES - 90% equities - 10% options

LONG WEIGHTAGE & SHORT WEIGHTAGE - 100% & 75% of allocated stock capital

Basic Portfolio & Risk Management for similar strategy done here by other user, dont know what might be helpful  - https://www.quantconnect.com/terminal/index.php?key=processCache&request=embedded_backtest_692142e204cbcc6483a9c28f8d452cb4.html#


'''




import datetime

class SuperTrendTester(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 7, 1)  # Start Date
        #self.SetEndDate(2021, 1, 15)  # End Date
        
        self.SetCash(100000)  # Set Strategy Cash
        self.spy = 'SPY'
        self.equity = self.AddEquity('SPY', Resolution.Minute)
        BarPeriod = TimeSpan.FromMinutes(5)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.closeWindow = []
        self.highWindow = []
        self.lowWindow = []
        self.atrWindow = []
        self.atrDown = []
        self.atrUp = []
        self.ha = self.HeikinAshi(self.spy, Resolution.Minute)
        self.atr = self.ATR(self.spy, 200, Resolution.Minute)
        #self.heikinAshi = SymbolData(self, self.spy)
        #self.superTrend = SuperTrend(self, self.spy, 2, 7, datetime.timedelta(minutes=5))
        consolidator = TradeBarConsolidator(timedelta(minutes=5))
        consolidator.DataConsolidated += self.OnDataConsolidated
        self.SubscriptionManager.AddConsolidator(self.spy, consolidator)
        self.RegisterIndicator(self.spy, self.ha,  consolidator)
        self.RegisterIndicator(self.spy, self.ha,  consolidator)
        
        self.putsheld = False
        self.callsheld = False
        self.noneheld = True
        
        self.underlyingsymbol = 'SPY'
        self.syl = 'SPY'

    def OnDataConsolidated(self, sender, data):
        
        if not self.ha.IsReady: return
        if not self.atr.IsReady: return
    
        self.closeWindow.insert(0,self.ha.Close.Current.Value)
        if len(self.closeWindow) > 2:
            self.closeWindow.pop()
        self.highWindow.insert(0,self.ha.High.Current.Value)
        if len(self.highWindow) > 2:
            self.highWindow.pop()
        self.lowWindow.insert(0,self.ha.Low.Current.Value)
        if len(self.lowWindow) > 2:
            self.lowWindow.pop()
        self.atrWindow.insert(0,self.atr.Current.Value)
        if len(self.atrWindow) > 2:
            self.atrWindow.pop()
        self.multiplier = 1
        self.Value = None
        try:
            hl = (self.highWindow[0] + self.lowWindow[0]) / 2
            hltwo = (self.highWindow[0] + self.lowWindow[0]) / 2
            hltwoPrev = (self.highWindow[1] + self.lowWindow[1]) / 2
    
            downNow = hltwo - self.multiplier * self.atrWindow[0]
            downPrev = hltwoPrev - self.multiplier * self.atrWindow[1]
            atrDown = max(downPrev, downNow) if self.closeWindow[1] > downPrev else downNow
            self.atrDown.insert(0,atrDown)
    
            upNow = hltwo + self.multiplier * self.atrWindow[0]
            upPrev = hltwoPrev + self.multiplier * self.atrWindow[1]
            atrUp = min(upNow, upPrev) if self.closeWindow[1] < upPrev else upNow
            self.atrUp.insert(0,atrUp)
        except:
            return
        
        try:
            if self.closeWindow[0] > self.atrUp[1]:
                self.Value = self.atrDown[0]
            elif self.closeWindow[0] < self.atrDown[1]:
                self.Value = self.atrUp[0]
            else:
                pass
            if self.Securities['SPY'].Price < self.Value:# and not self.putsheld:
                self.BuyPuts()
                self.SetHoldings('SPY', 1)
                    
            if self.Securities['SPY'].Price > self.Value:# and not self.callsheld:
                self.BuyCalls()
                self.SetHoldings('SPY', -1)
        except:
            return
        
        

    def OnData(self, data):
        pass
        #self.Plot("Chart", "SPY", data[self.spy].Close)
        #if self.superTrend.Value is None: return
        #self.Plot("Chart", "SuperTrend", self.superTrend.Value)
        
        #expiries = [x.Key.ID.OptionRight for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        #if not self.Portfolio.Invested:
    
    def BuyPuts(self):
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
        if len(contracts) == 0: return
        filtered_contracts = self.InitialFilter(self.underlyingsymbol, contracts, -12, -5, 0, 7)
        put = [x for x in filtered_contracts if x.ID.OptionRight == 1] 
        # sorted the contracts according to their expiration dates and choose the ATM options
        contracts = sorted(sorted(put, key = lambda x: abs(self.Securities[self.syl].Price - x.ID.StrikePrice)), 
                                        key = lambda x: x.ID.Date, reverse=True)
        self.contract = contracts[0]
        self.AddOptionContract(self.contract, Resolution.Minute)
        self.Debug(str(self.Time) + str(self.contract) + "Put")
        self.Buy(self.contract, 10)
        self.callsheld = False
        self.putsheld = True
        
    def BuyCalls(self):
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
        if len(contracts) == 0: return
        filtered_contracts = self.InitialFilter(self.underlyingsymbol, contracts, +5, +12, 0, 7)
        call = [x for x in filtered_contracts if x.ID.OptionRight == 0] 
        for a in call:
            abc = a
        # sorted the contracts according to their expiration dates and choose the ATM options
        contracts = sorted(sorted(call, key = lambda x: abs(self.Securities[self.syl].Price - x.ID.StrikePrice)), 
                                        key = lambda x: x.ID.Date, reverse=True)
        self.contract = contracts[10]
        self.AddOptionContract(self.contract, Resolution.Minute)
        self.Debug(str(self.Time) + str(self.contract) + "Call")
        self.Buy(self.contract, 10)
        self.callsheld = True
        self.putsheld = False
        
    def InitialFilter(self, underlyingsymbol, symbol_list, min_strike_rank, max_strike_rank, min_expiry, max_expiry):
            
        if len(symbol_list) == 0 : return
        # fitler the contracts based on the expiry range
        contract_list = [i for i in symbol_list if min_expiry < (i.ID.Date.date() - self.Time.date()).days < max_expiry]
        # find the strike price of ATM option
        atm_strike = sorted(contract_list,
                            key = lambda x: abs(x.ID.StrikePrice - self.Securities[underlyingsymbol].Price))[0].ID.StrikePrice
        strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
        # find the index of ATM strike in the sorted strike list
        atm_strike_rank = strike_list.index(atm_strike)
        try: 
            min_strike = strike_list[atm_strike_rank + min_strike_rank]
            max_strike = strike_list[atm_strike_rank + max_strike_rank]
        except:
            min_strike = strike_list[0]
            max_strike = strike_list[-1]
           
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike and i.ID.StrikePrice <= max_strike]

        return filtered_contracts