# https://quantpedia.com/strategies/dispersion-trading/
#
# The investment universe consists of stocks from the S&P 100 index. Trading vehicles are options on stocks from this index and also options on the index itself. The investor uses analyst forecasts of earnings per share
# from the Institutional Brokers Estimate System (I/B/E/S) database and computes for each firm the mean absolute difference scaled by an indicator of earnings uncertainty (see page 24 in the source academic paper for 
# detailed methodology). Each month, investor sorts stocks into quintiles based on the size of belief disagreement. He buys puts of stocks with the highest belief disagreement and sells the index puts with Black-Scholes 
# deltas ranging from -0.8 to -0.2.
#
# QC Implementation changes:
#   - Due to lack of data, strategy only buys puts of 100 liquid US stocks and sells the SPX index puts.

#region imports
from AlgorithmImports import *
from universe import IndexConstituentsUniverseSelectionModel
from numpy import floor
#endregion

class DispersionTrading(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2010, 1, 1)
        self.SetCash(1_000_000)
        
        self.min_expiry:int = 20
        self.max_expiry:int = 60

        self.buying_power_model:int = 2
        
        self._etf: str = 'SPY'
        self._market_cap_count: int = 100
        self._last_selection: List[Symbol] = []
        self._subscribed_contracts = {}
        self._price_data: Dict[Symbol, RollingWindow] = {}
        self._period: int = 21

        self.index_symbol:Symbol = self.AddIndex('SPX', Resolution.DAILY).Symbol
        self.percentage_traded:float = 1.0
        
        self._last_selection:List[Symbol] = []

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.daily_precise_end_time = False
        self.set_security_initializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.RAW))
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.schedule.on(self.date_rules.month_start())
        self.universe_settings.resolution = Resolution.DAILY

        self.add_universe_selection(
            IndexConstituentsUniverseSelectionModel(
                self._etf, 
                self._market_cap_count, 
                self.universe_settings
            )
        )

        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.
        self.settings.daily_precise_end_time = False

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        # update index constituents
        for security in changes.added_securities:
            if security.subscriptions[0].security_type != SecurityType.EQUITY:
                continue
            symbol: Symbol = security.symbol
            if symbol == self.index_symbol:
                continue
            self._last_selection.append(security.symbol)

        # for security in changes.removed_securities:
        #     if security.symbol in self._last_selection:
        #         self._last_selection.remove(security.symbol)

    def OnData(self, data: Slice) -> None:
        # liquidate portfolio, when SPX contract is about to expire in 2 days
        if self.index_symbol in self._subscribed_contracts and self._subscribed_contracts[self.index_symbol].ID.date.date() - timedelta(2) <= self.time.date():
            self._subscribed_contracts.clear()   # perform new subscribtion
            self.liquidate()
            
        if len(self._subscribed_contracts) == 0:
            if self.portfolio.invested:
                self.liquidate()
            
            # NOTE order is important, index should come first
            for symbol in [self.index_symbol] + self._last_selection:
                if symbol != self.index_symbol:
                    if symbol not in data:
                        continue
                if self.Securities[symbol].IsDelisted:
                    continue

                # subscribe to contract
                chain: OptionChain = self.option_chain(symbol)
                contracts: List[OptionContract] = [i for i in chain]

                if len(contracts) == 0:
                    continue

                # get current price for stock
                underlying_price:float = self.securities[symbol].ask_price
                
                # get strikes from stock contracts
                strikes: List[float] = [i.strike for i in contracts]
                
                # check if there is at least one strike    
                if len(strikes) <= 0:
                    continue
            
                # at the money
                atm_strike: float = min(strikes, key=lambda x: abs(x-underlying_price))

                atm_puts: Symbol = sorted(
                    filter(
                        lambda x: x.right == OptionRight.PUT 
                        and x.strike == atm_strike
                        and self.min_expiry <= (x.expiry - self.Time).days <= self.max_expiry, 
                        contracts
                    ), 
                    key=lambda item: item.expiry, reverse=True
                )

                # index contract is found
                if symbol == self.index_symbol and len(atm_puts) == 0:
                    # cancel whole selection since index contract was not found
                    return
                    
                # make sure there are enough contracts
                if len(atm_puts) > 0:
                    # add contract
                    option = self.AddOptionContract(atm_puts[0], Resolution.DAILY)
                    option.PriceModel = OptionPriceModels.CrankNicolsonFD()
                    option.SetDataNormalizationMode(DataNormalizationMode.Raw)

                    # store subscribed atm put contract
                    self._subscribed_contracts[symbol] = atm_puts[0].symbol
        
        if self.index_symbol not in self._subscribed_contracts:
            self._subscribed_contracts.clear()
            self._last_selection.clear()
            return

        # perform trade, when spx and stocks contracts are selected            
        if (not self.Portfolio.Invested and len(self._subscribed_contracts) != 0 and self.index_symbol in self._subscribed_contracts):
            index_option_contract = self._subscribed_contracts[self.index_symbol]
            # make sure subscribed SPX contract has data
            if self.Securities.ContainsKey(index_option_contract):
                if self.Securities[index_option_contract].Price != 0 and self.Securities[index_option_contract].IsTradable:
                    # sell SPX ATM put contract
                    self.Securities[index_option_contract].MarginModel = BuyingPowerModel(self.buying_power_model)
                    price:float = self.Securities[self.index_symbol].ask_price
                    if price != 0:
                        if index_option_contract.value == 'SPX   140419P01840000':  #TODO: Bug. Once opened position, cannot close it.
                            return
                        notional_value: float = (price * self.securities[index_option_contract].symbol_properties.contract_multiplier)
                        if notional_value != 0:
                            q: int = self.portfolio.total_portfolio_value * self.percentage_traded  // notional_value
                            self.sell(index_option_contract, q)
                            # self.market_order(index_option_contract, -q)

                    # buy stock's ATM put contracts            
                    long_count:int = len(self._subscribed_contracts) - 1     # minus index symbol
                    for stock_symbol, stock_option_contract in self._subscribed_contracts.items():
                        if stock_symbol == self.index_symbol:
                            continue
                        
                        if stock_option_contract in data and data[stock_option_contract]:
                            if self.Securities[stock_option_contract].Price != 0 and self.Securities[stock_option_contract].IsTradable:
                                # buy contract
                                self.Securities[stock_option_contract].MarginModel = BuyingPowerModel(self.buying_power_model)
                                if self.Securities.ContainsKey(stock_option_contract):
                                    price:float = self.Securities[stock_symbol].ask_price
                                    if price != 0:
                                        notional_value: float = (price * self.securities[stock_option_contract].symbol_properties.contract_multiplier)
                                        if notional_value != 0:
                                            q: int = self.portfolio.total_portfolio_value * self.percentage_traded // long_count // notional_value
                                            self.buy(stock_option_contract, q)
                                            # self.market_order(stock_option_contract, q)

# Custom fee model
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))

#region imports
from AlgorithmImports import *
#endregion

class IndexConstituentsUniverseSelectionModel(ETFConstituentsUniverseSelectionModel):
    def __init__(
        self, 
        etf: str, 
        top_market_cap_count: int,
        universe_settings: UniverseSettings = None
    ) -> None:

        symbol = Symbol.create(etf, SecurityType.EQUITY, Market.USA)
        self._top_market_cap_count: int = top_market_cap_count
        
        super().__init__(
            symbol, 
            universe_settings,
            universe_filter_func=self._etf_constituents_filter
        )

    def _etf_constituents_filter(self, constituents: List[ETFConstituentUniverse]) -> List[Symbol]:
        # select n largest equities in the ETF
        selected = sorted(
            [c for c in constituents if c.weight],
            key=lambda c: c.weight, reverse=True
        )[:self._top_market_cap_count]

        return list(map(lambda x: x.symbol, selected))