from SymbolData import *


class RetrospectiveFluorescentPinkWhale(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 11, 26)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        # ~~~ static vs dyanmic equity universe ~~~
        
        # Static Universe
        # 
        # OptionChainProvider vs AddOption
        # If you're using greeks, you'll need to use AddOption
        # 
        # # AddOption leads to slow backtests, basically impossible to use with dynamic
        # 
        # OptionChainProvider - faster backtests but lack of greeks and IV
        
        # If no greeks/IV needed -> OptionChainProvider
        
        # Dynamic Equity Universe
        # OptionChainProvider leads to much faster backtests
        #
    
        # If using AddOption, generally will use bottom technique to retrieve contracts
        # data.OptionsChains[option_chain.Symbol(identifier for chain)].Contracts # list of OptionContract objects
        
        # OptionContract.Greeks.Delta, etc
        
        
        # Medium sized static universe
        
        self.benchmark = "SPY"
        
        tickers = ["SPY", "QQQ", "GLD", "TLT"]
        
        self.symbol_data = {}
        
        for ticker in tickers:
            equity = self.AddEquity(ticker, Resolution.Minute)
            # this will be done automatically 
            equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        
        # rebalance scheduled event
        self.Schedule.On(self.DateRules.EveryDay(self.benchmark), self.TimeRules.BeforeMarketClose(self.benchmark, 5), self.Rebalance)
        
        
        self.SetSecurityInitializer(lambda security: security.SetMarketPrice(self.GetLastKnownPrice(security)))
        
        
        ### STRATEGY ####
        # Entry: 5 minutes before close daily (35 day Current MA) > (35 day MA - 5 periods) -> Long
        # Write Weekly ATM put, and buy ~ $10 below market price
        
        # Just let it expire with either max profit or close before expiration, buy back for loss
    
        # open_option_positions = [symbol for symbol in self.Portfolio.Keys if self.Portfolio[symbol].Invested and \
        #     symbol.SecurityType == SecurityType.Option]
            
        # # 30 mins before market close on friday/OEC
        # for symbol in open_option_positions:
        #     # close out
        #     self.Liquidate(symbol)
    
    
    
    
    def Rebalance(self):
        
        
        for symbol, symbol_data in self.symbol_data.items():
                
            if symbol.Value != "SPY":
                continue;
            
            if not symbol_data.sma_window.IsReady:
                continue
            signal = symbol_data.long_signal
            
            if signal and not self.Portfolio[symbol].Invested:
                # day of week - self.Time.weekday(), from there calculate dte needed
                
                current_market_price = self.Securities[symbol].Price
                short_put = self.select_option(symbol, OptionRight.Put, current_market_price, 7)
                long_put = self.select_option(symbol, OptionRight.Put, current_market_price - 25, 7)
                
                # could put this inside of select_option
                
                # nuanced detail - when we call AddOptionContract to subscribe to
                # option data, data is not immediately available, it takes 1 time step
                # so if we try to call market order/ limit order, it won't be able to fill
                
                self.AddOptionContract(short_put)
                self.AddOptionContract(long_put)
                
                short_put_bid = self.Securities[short_put].BidPrice
                long_put_ask = self.Securities[short_put].AskPrice
                
                
                self.LimitOrder(short_put, -1,  short_put_bid)
                self.LimitOrder(long_put, 1, long_put_ask)
                
                
    
    def select_option(self, underlying_symbol, option_right, strike, expiry_dte):
        
        # underlying symbol and date of option data needed
        
        # looks into option open interest data for that underlying
        # returns that
        
        # list of Symbol objects
        contract_symbols = self.OptionChainProvider.GetOptionContractList(underlying_symbol, self.Time)
        
        
        filtered_for_right = [symbol for symbol in contract_symbols if symbol.ID.OptionRight == option_right]
        
        
        expiry_date = self.Time + timedelta(expiry_dte)
        
        
        # ascending order - reverse=False - default
        sorted_by_expiry = sorted(filtered_for_right, key=lambda s:abs(s.ID.Date - expiry_date), reverse=False)
        
        
        if len(sorted_by_expiry) > 0:
            desired_expiry_date = sorted_by_expiry[0].ID.Date
        else:
            self.Debug(f"No contracts found for {underlying_symbol}")
            return None
        
        
        contracts_with_desired_expiry = [symbol for symbol in filtered_for_right if symbol.ID.Date == desired_expiry_date]
        
        # sorted ascending again    
        sorted_by_strike = sorted(contracts_with_desired_expiry, key=lambda s: abs(s.ID.StrikePrice - strike), reverse=True)        
        
        
        if len(sorted_by_strike) > 0:
            # returns a symbol object for selected option
            return sorted_by_strike[0]
        else:
            return None
        
        
        
        
        
        
        
        

    
    def OnSecuritiesChanged(self, changes):
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            
            if not symbol in self.symbol_data:
                self.symbol_data[symbol] = SymbolData(self, symbol)
        
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            
            if symbol in self.symbol_data:
                symbol_data = self.symbol_data.pop(symbol, None)
                symbol_data.kill_data_subscriptions()
    
    
    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
        '''

        

class SymbolData:
    
    def __init__(self, algorithm, symbol):
        
        
        self.algorithm = algorithm
        self.symbol = symbol
        
        # self.algorithm.SMA
        self.sma = SimpleMovingAverage(35)
        
        # User defined method OnSMA will fire everytime SMA is updated
        self.sma.Updated += self.OnSMA
        
        # consolidator is subscribed to receive data
        self.daily_consolidator = QuoteBarConsolidator(timedelta(days=1))
        self.algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.daily_consolidator)
        
        # consolidator will update to indicator
        self.algorithm.RegisterIndicator(self.symbol, self.sma, self.daily_consolidator)
        
        self.sma_window = RollingWindow[IndicatorDataPoint](5)
        
        # pandas dataframe stuff is buggy
        # self.warm_up_sma()
        
        
    def warm_up_sma(self):
        
        # option 1: update consolidator with minute historical data -> spit out daily bar -> update sma
        # option 2: daily historical data -> update sma
        
        # daily data in
        # 0th index is symbol, 1st index is Time
        history = self.algorithm.History(self.symbol, 35, Resolution.Daily).unstack(0).close
        
        # time series daily close data, with each column being symbol
        if not history.empty:
            
            if str(self.symbol) in history.columns:
                
                close_pandas_series_data = history[str(self.symbol)]
                
                for time, close in close_pandas_series_data:
                    
                    self.sma.Update(time, close)
    
    def OnSMA(self, sender, updated):
        
        if self.sma.IsReady:
            
            # self.sma - SimpleMovingAverage/Indicator Object
            # self.sma.Current = IndicatorDataPoint object
            # self.sma.Current.Value - float object
            
            self.sma_window.Add(self.sma.Current)
    
    
    
    def kill_data_subscriptions(self):
        
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.daily_consolidator)
        
    
    @property
    def long_signal(self):
        
        return self.sma.Current.Value > self.sma_window[4].Value