# region imports
from AlgorithmImports import *
# endregion


class EarningsReversal(QCAlgorithm):

    def Initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        # ------------- Paramters --------------
        # Minimum percentage decline before entry
        self._entry_move = 0.02
        # Max number of positions allowed at once (equal weighting)
        self._max_positions = 10
        # Size of first universe selection liquidity filter
        self._liquidity_filter_size = 6*self._max_positions
        # Number of days past since earnings
        self._days_since_earnings = 1
        # Percentage offset of trailing stop loss
        self._stop_loss = 0.10
        # --------------------------------------
        # Add a universe of US Equities.
        self._universe = self.add_universe(self._select_assets)
        # Add a Scheduled Event to open new positions.
        self.schedule.on(
            self.date_rules.every_day('SPY'),
            self.time_rules.at(8, 0),
            self._trade
        )
        # Add a custom bar chart.
        chart = Chart('Positions')
        chart.add_series(Series('Longs', SeriesType.BAR, 0))
        self.add_chart(chart)

    def _select_assets(self, fundamentals):
        # Select the most liquid stocks trading above $5.
        fundamentals = [f for f in fundamentals if f.price > 5 and f.has_fundamental_data]
        fundamentals = sorted(fundamentals, key=lambda f: f.dollar_volume, reverse=True)[:self._liquidity_filter_size]
        # Select the subset of stocks with recent earnings.
        fundamentals = [
            f for f in fundamentals 
            if self.time == f.earning_reports.file_date.value + timedelta(self._days_since_earnings)
        ]
        # Select the subset of stocks that have fallen in price since the earnings release.
        prices_around_earnings = self.history([f.symbol for f in fundamentals], self._days_since_earnings+3, Resolution.DAILY)
        universe = []
        for f in fundamentals:
            # Find tradeable date closest to specified number of days before earnings
            date = min(
                prices_around_earnings.loc[f.symbol]["close"].index, 
                key=lambda x:abs(x-(f.earning_reports.file_date.value - timedelta(1)))
            )
            price_on_earnings = prices_around_earnings.loc[f.symbol]["close"][date]
            # Check if stock fell far enough.
            if price_on_earnings * (1-self._entry_move) > f.price:
                universe.append(f.symbol)
        return universe
    
    # Open equal weighted positions with trailing stop losses.
    def _trade(self):
        # Plot number of existing positions
        positions = [symbol for symbol, holding in self.portfolio.items() if holding.invested]
        self.plot("Positions", "Longs", len(positions))
        # Enter new trades if there are assets in the universe.
        if not self._universe.selected:
            return
        available_trades = self._max_positions - len(positions)
        for symbol in [x for x in self._universe.selected if x not in positions][:available_trades]:
            # Buy the stock.
            self.set_holdings(symbol, 1 / self._max_positions)
        
    def on_order_event(self, order_event: OrderEvent) -> None:
        if (order_event.status != OrderStatus.FILLED or 
            order_event.ticket.order_type != OrderType.MARKET_ON_OPEN):
            return
        # Add a trailing stop order.
        self.trailing_stop_order(order_event.symbol, -order_event.quantity, self._stop_loss, True)