from AlgorithmImports import *
import numpy as np

class StandardizedUnexpectedEarnings(QCAlgorithm):
    '''Step 1. Calculate the change in quarterly EPS from its value four quarters ago
       Step 2. Calculate the st dev of this change over the prior eight quarters
       Step 3. Get standardized unexpected earnings (SUE) from dividing results of step 1 by step 2
       Step 4. Each month, sort universe by SUE and long the top quantile
       Reference:
       [1] Foster, Olsen and Shevlin, 1984, Earnings Releases, Anomalies, and the Behavior of Security Returns,
           The Accounting Review.
       [2] Hou, Xue and Zhang, 2018, Replicating Anomalies, Review of Financial Studies.
    '''

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100000)
        self.settings.seed_initial_prices = True        
        # Define some parameters.
        self._months_eps_change = 12
        self._months_count = 36
        self._num_selected = 1000
        self._top_percent = 0.05        
        self._eps_by_symbol = {}
        # Add a universe of US Equities.
        self._date_rule = self.date_rules.month_start("SPY") 
        self.universe_settings.schedule.on(self._date_rule)
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select)        
        self.set_warm_up(timedelta(self._months_count*31+30))
    
    def _select(self, fundamentals):
        # Filter to stocks with valid price and EPS data, select top by dollar volume.
        selected = sorted(
            [f for f in fundamentals if f.price > 5 and f.earning_reports.basic_eps.three_months],
            key=lambda f: f.dollar_volume
        )[-self._num_selected:]        
        # Update EPS rolling windows and calculate SUE for stocks with sufficient history.
        sue_by_symbol = {}
        for f in selected:
            # Update the trailing EPS data.
            if f.symbol not in self._eps_by_symbol:
                self._eps_by_symbol[f.symbol] = RollingWindow[float](self._months_count)
            eps_window = self._eps_by_symbol[f.symbol]
            eps_window.add(f.earning_reports.basic_eps.three_months)
            # Calculate SUE if there is enough data.
            if not eps_window.is_ready:
                continue
            # Current quarter EPS change from four quarters ago.
            eps_change = eps_window[0] - eps_window[self._months_eps_change]            
            # Calculate historical EPS changes for standard deviation.
            new_eps = list(eps_window)[:self._months_count - self._months_eps_change:3]
            old_eps = list(eps_window)[self._months_eps_change::3]
            changes = [float(new - old) for new, old in zip(new_eps, old_eps)]
            std = np.std(changes)
            if std > 0:
                sue_by_symbol[f.symbol] = eps_change / std
        # Select top quantile by SUE.
        if not sue_by_symbol:
            return []
        num_long = max(1, int(self._top_percent * len(sue_by_symbol)))
        return sorted(sue_by_symbol, key=lambda symbol: sue_by_symbol[symbol])[-num_long:]
        
    def on_warmup_finished(self):
        # Add a Scheduled event to rebalance the portfolio monthly.
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self._date_rule, time_rule, self._rebalance)
        # Rebalance the portfolio today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _rebalance(self):
        # Form an equal-weighted portfolio..
        weight = 1.0 / len(self._universe.selected)
        self.set_holdings([PortfolioTarget(symbol, weight) for symbol in self._universe.selected], True)