# region imports
from AlgorithmImports import *
from dataclasses import dataclass
# endregion

# NOTE: Manager for new trades. It's represented by certain count of equally weighted brackets for long and short positions.
# If there's a place for new trade, it will be managed for time of holding period.
class TradeManager():
    def __init__(self, 
                algorithm: QCAlgorithm, 
                long_size: int, 
                short_size: int, 
                holding_period: int,) -> None:
        self.algorithm: QCAlgorithm = algorithm  # algorithm to execute orders in.
        
        self.long_size: int = long_size
        self.short_size: int = short_size
        
        self.long_len: int = 0
        self.short_len: int = 0

        # Arrays of ManagedSymbols
        self.symbols: List[ManagedSymbol] = []
        
        self.holding_period: int = holding_period    # Days of holding.
    
    # Add stock symbol object
    def add(
        self, 
        symbol: Symbol, 
        long_flag: bool, 
        price: float) -> None:
    
        # Open new long trade.
        quantity: Union[None, int] = None
        
        if long_flag:
            # If there's a place for it.
            if self.long_len < self.long_size:
                quantity: int = int(self.algorithm.portfolio.total_portfolio_value / self.long_size / price)

                self.algorithm.market_order(symbol, quantity)
                self.long_len += 1
            # else:
            #     self.algorithm.Log("There's not place for additional trade.")

        # Open new short trade.
        else:
            # If there's a place for it.
            if self.short_len < self.short_size:
                quantity: int = int(self.algorithm.portfolio.total_portfolio_value / self.short_size / price)

                self.algorithm.market_order(symbol, -quantity)
                self.short_len += 1
            # else:
                # self.algorithm.Log("There's not place for additional trade.")
   
        if quantity:
            managed_symbol: ManagedSymbol = ManagedSymbol(symbol, self.holding_period, quantity, long_flag)
            self.symbols.append(managed_symbol)

    # Decrement holding period and liquidate symbols.
    def try_liquidate(self) -> None:
        symbols_to_delete: List[ManagedSymbol] = []
        for managed_symbol in self.symbols:
            managed_symbol.days_to_liquidate -= 1
            
            # Liquidate.
            if managed_symbol.days_to_liquidate == 0:
                symbols_to_delete.append(managed_symbol)
                    
                if managed_symbol.long_flag:
                    self.algorithm.market_order(managed_symbol.symbol, -managed_symbol.quantity)
                    self.long_len -= 1
                else:
                    self.algorithm.market_order(managed_symbol.symbol, managed_symbol.quantity)
                    self.short_len -= 1
                

        # Remove symbols from management.
        for managed_symbol in symbols_to_delete:
            self.symbols.remove(managed_symbol)
    
    def liquidate_ticker(self, ticker: str) -> None:
        symbol_to_delete: Union[None, Symbol] = None
        for managed_symbol in self.symbols:
            if managed_symbol.symbol.Value == ticker:
                if managed_symbol.long_flag:
                    self.algorithm.market_order(managed_symbol.symbol, -managed_symbol.quantity)
                    self.long_len -= 1
                else:
                    self.short_len -= 1
                    self.algorithm.market_order(managed_symbol.symbol, managed_symbol.quantity)
                symbol_to_delete = managed_symbol
                
                break
        
        if symbol_to_delete: self.symbols.remove(symbol_to_delete)
        else: self.algorithm.Debug("Ticker is not held in portfolio!")

@dataclass   
class ManagedSymbol():
    symbol: Symbol 
    days_to_liquidate: int 
    quantity: int 
    long_flag: bool

# https://quantpedia.com/strategies/ipos-negative-returns-after-options-listing
# 
# The sample consists of U.S. companies that went public between 1996 and 2017 and had options listed within 156 weeks (3 years) after the IPO issuance date 
# (newly public companies, specifically those that have gone public within the last three years and have options listed on their equity). Individual instruments 
# are selected based on the presence of options trading activity and the affiliation of the IPO lead underwriter with the Options Clearing Corporation (OCC). 
# Companies where the lead underwriter is an OCC member are prioritized, as this affiliation suggests a potential private information source that proprietary 
# traders might exploit.
# (U.S. initial public offerings data on ordinary common shares from the Securities DataCompany (SDC) New Issues database, security prices, and trading 
# information from the Center for Research in
# Security Prices (CRSP), information on option prices and exchanges’ listing dates is collected from OptionMetrics, which covers all U.S. exchange-listed 
# equities and market indices, and security-lending data from Markit Securities Finance, open-close options volume data from 8 options exchange: acquired from 
# two sources: the Chicago Board of Options Exchange (CBOE) and the NASDAQ, each operating multiple options exchange.)
# Main Principle Recapitulation: The strategy idea is is capitalizing on finding that when options are listed for newly issued IPOs, their returns decline. In 
# ideal circumstances, investor should start shorting IPO stocks 12 weeks before options are listed for newly issued IPOs and cover short position 24 weeks 
# after options are introduced.
# Rebalancing & Weighting: Rebalanced dynamically, weekly, if there is a corresponding signal for inclusion (see first sentence of the first paragraph). 
# Consider weighting options in a value-weighted manner.
# 
# QC Implementation changes:
#   - Trading is executed on the date when the first option chain becomes available for the given stock.

# region imports
from AlgorithmImports import *
from dateutil.relativedelta import relativedelta
from pandas.tseries.offsets import BDay
import data_tools
# endregion

class IPOsNegativeReturnsAfterOptionsListing(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2013, 1, 1)
        self.set_cash(100_000)

        self._tickers_to_ignore: List[str] = ['CLSKW', 'JGGCR']
        self.exchange_codes: List[str] = ['NYS', 'NAS', 'ASE']    

        self._option_trade_flag: bool = True
        self._leverage: int = 5
        self._quantile: int = 5
        self._period: int = 12
        self._holding_period: int = 24  # weeks
        self._short_count: int = 50

        self._ipo_dates: Dict[Symbol, datetime] = {}
    
        self._trade_manager: data_tools.TradeManager = data_tools.TradeManager(self, 0, self._short_count, self._holding_period)

        market: Symbol = self.add_equity('SPY', Resolution.DAILY).Symbol

        self.universe_settings.leverage = self._leverage
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(EODHDUpcomingIPOs, self.ipo_selection_filter)
        self.settings.daily_precise_end_time = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.

        self.schedule.on(
            self.date_rules.week_start(market),
            self.time_rules.after_market_open(market),
            self.selection
        )

    def ipo_selection_filter(self, ipos: List[EODHDUpcomingIPOs]) -> List[Symbol]:
        # Save upcoming IPOs.
        for ipo_data in ipos:
            if ipo_data.symbol not in self._ipo_dates:
                if ipo_data.ipo_date and ipo_data.symbol.value not in self._tickers_to_ignore:
                    self._ipo_dates[ipo_data.symbol] = (ipo_data.ipo_date + BDay(1)).date()

        return list(self._ipo_dates.keys())

    def on_data(self, slice: Slice) -> None:
        symbols_to_delete: List[Symbol] = []

        # Trade execution.
        for symbol, date in self._ipo_dates.items():
            chain = self.option_chain(symbol, flatten=True).data_frame
            trade_flag: bool = not chain.empty if self._option_trade_flag else self._ipo_dates[symbol] == self.time.date()
            if trade_flag:
                if date > self.time.date():
                    continue
                if slice.contains_key(symbol) and slice[symbol]:
                    self._trade_manager.add(symbol, False, slice[symbol].close)
                    symbols_to_delete.append(symbol)

        # Remove traded stocks.
        for symbol in symbols_to_delete:
            self._ipo_dates.pop(symbol)
        symbols_to_delete.clear()

    def selection(self) -> None:
        # Hold for n weeks.
        self._trade_manager.try_liquidate()