# https://quantpedia.com/strategies/post-earnings-announcement-effect/
#
# The investment universe consists of all stocks from NYSE, AMEX, and NASDAQ except financial and utility firms and stocks with prices less than $5.
# Two factors are used: EAR (Earnings Announcement Return) and SUE (Standardized Unexpected Earnings). SUE is constructed by dividing the earnings 
# surprise (calculated as actual earnings minus expected earnings; expected earnings are computed using a seasonal random walk model with drift)
# by the standard deviation of earnings surprises. EAR is the abnormal return for firms recorded over a three-day window centered on the last 
# announcement date, in excess of the return of a portfolio of firms with similar risk exposures.
# Stocks are sorted into quantiles based on the EAR and SUE. To avoid look-ahead bias, data from the previous quarter are used to sort stocks. 
# Stocks are weighted equally in each quantile. The investor goes long stocks from the intersection of top SUE and EAR quantiles and goes short
# stocks from the intersection of the bottom SUE and EAR quantiles the second day after the actual earnings announcement and holds the portfolio 
# one quarter (or 60 working days). The portfolio is rebalanced every quarter.
#
# QC Implementation changes:
#   - Universe consists of stocks, with earnings data from https://www.nasdaq.com/market-activity/earnings available.
#   - At least 4 years of seasonal earnings data is required to calculate earnigns surprise.
#   - At least 4 years of earnings surprise values are required for SUE calculation.
#   - Only long leg is traded since research paper claims that major part of strategy performance is produced by long leg only.

#region imports
from AlgorithmImports import *
import numpy as np
from collections import deque
from pandas.tseries.offsets import BDay
from dateutil.relativedelta import relativedelta
from typing import Dict, List, Tuple, Deque
#endregion

class PostEarningsAnnouncementEffect(QCAlgorithm):

    def Initialize(self) -> None:
        self.SetStartDate(2010, 1, 1)
        self.SetCash(100_000)

        self.earnings_surprise: Dict[Symbol, float] = {}
        self.min_seasonal_eps_period: int = 4
        self.min_surprise_period: int = 4
        self.leverage: int = 5
        self.percentile_range: List[int] = [80, 20]
        
        self.long: List[Symbol] = []
        
        # SUE and EAR history for previous quarter used for statistics.
        self.sue_ear_history_previous: List[Tuple[float, float]] = []
        self.sue_ear_history_actual: List[Tuple[float, float]] = []
        
        # EPS data keyed by tickers, which are keyed by dates
        self.eps_by_ticker: Dict[str, float] = {}
        
        # daily price data
        self.price_data_with_date: Dict[Symbol, Deque[float]] = {}
        self.price_period: int = 63

        self.market: Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.price_data_with_date[self.market] = deque(maxlen=self.price_period)

        # parse earnings dataset
        self.first_date: Union[datetime.date, None] = None
        earnings_data: str = self.Download('data.quantpedia.com/backtesting_data/economic/earnings_dates_eps.json')
        earnings_data_json: List[dict] = json.loads(earnings_data)
        
        for obj in earnings_data_json:
            date: datetime.date = datetime.strptime(obj['date'], "%Y-%m-%d").date()
            
            if not self.first_date: self.first_date = date

            for stock_data in obj['stocks']:
                ticker: str = stock_data['ticker']

                if stock_data['eps'] == '':
                    continue

                # initialize dictionary for dates for specific ticker
                if ticker not in self.eps_by_ticker:
                    self.eps_by_ticker[ticker] = {}
                
                # store EPS value keyed date, which is keyed by ticker
                self.eps_by_ticker[ticker][date] = float(stock_data['eps'])

        self.month: int = 12
        self.selection_flag: bool = False
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.
        self.settings.daily_precise_end_time = False
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.FundamentalSelectionFunction)
        self.Schedule.On(self.DateRules.MonthStart(self.market), self.TimeRules.AfterMarketOpen(self.market), self.Selection)

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(self.leverage)
        
        # remove earnings surprise data so it remains consecutive
        for security in changes.RemovedSecurities:
            symbol: Symbol = security.Symbol
            if symbol in self.earnings_surprise:
                del self.earnings_surprise[symbol]
    
    def FundamentalSelectionFunction(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # update daily price data
        for stock in fundamental:
            symbol: Symbol = stock.Symbol

            if symbol in self.price_data_with_date:
                self.price_data_with_date[symbol].append((self.Time.date(), stock.AdjustedPrice))
        
        if not self.selection_flag:
            return Universe.Unchanged
        self.selection_flag = False
        
        # filter only symbols, which have earnings data from csv
        selected: List[Symbol] = [
            x.Symbol for x in fundamental if x.Symbol.Value in self.eps_by_ticker
        ]
        
        # SUE and EAR data
        sue_ear: Dict[Symbol, float] = {}
        
        current_date: datetime.date = self.Time.date()
        prev_three_months: datetime = current_date - relativedelta(months=3)
        
        # warmup price data
        for symbol in selected:
            ticker: str = symbol.Value
            recent_eps_data: Union[None, datetime.date] = None

            if symbol not in self.price_data_with_date:
                self.price_data_with_date[symbol] = deque(maxlen=self.price_period)
                history: DataFrame = self.History(symbol, self.price_period, Resolution.Daily)
                if history.empty:
                    self.Log(f"Not enough data for {symbol} yet.")
                    continue
                closes: Series = history.loc[symbol].close
                for time, close in closes.items():
                    self.price_data_with_date[symbol].append((time.date(), close))
        
            # market price data is not ready yet
            if len(self.price_data_with_date[self.market]) != self.price_data_with_date[self.market].maxlen:
                return Universe.Unchanged

            if len(self.price_data_with_date[symbol]) != self.price_data_with_date[symbol].maxlen:
                continue 

            # store all EPS data since previous three months window    
            for date in self.eps_by_ticker[ticker]:
                if date < current_date and date >= prev_three_months:
                    EPS_value: float = self.eps_by_ticker[ticker][date]
                    
                    # create tuple (EPS date, EPS value of specific stock)
                    recent_eps_data: Tuple[datetime.date, float] = (date, EPS_value)
                    break
            
            if recent_eps_data:
                last_earnings_date: datetime.date = recent_eps_data[0]
                
                # get earnings history until previous earnings
                earnings_eps_history: List[Tuple[datetime.date, float]] = [
                    (x, self.eps_by_ticker[ticker][x]) for x in self.eps_by_ticker[ticker] if x < last_earnings_date
                ]
                
                # seasonal earnings for previous years
                # prev_month_date = last_earnings_date - relativedelta(months=1)
                # next_month_date = last_earnings_date + relativedelta(months=1)
                # month_range = [prev_month_date.month, last_earnings_date.month, next_month_date.month]
                
                # seasonal_eps_data = [x for x in earnings_eps_history if x[0].month in month_range]
                seasonal_eps_data: List[Tuple[datetime.date, float]] = [
                    x for x in earnings_eps_history if x[0].month == last_earnings_date.month
                ]
                
                if len(seasonal_eps_data) >= self.min_seasonal_eps_period:
                    # make sure we have a consecutive seasonal data. Same months with one year difference
                    year_diff: np.ndarray = np.diff([x[0].year for x in seasonal_eps_data])
                    if all(x == 1 for x in year_diff):
                        # SUE calculation
                        seasonal_eps: List[float] = [x[1] for x in seasonal_eps_data]
                        diff_values: np.ndarray = np.diff(seasonal_eps)
                        drift: float = np.average(diff_values)
                        
                        last_earnings_eps: float = seasonal_eps[-1]
                        expected_earnings: float = last_earnings_eps + drift
                        actual_earnings: float = recent_eps_data[1]
                        
                        earnings_surprise: float = actual_earnings - expected_earnings
                        
                        # initialize suprise data
                        if symbol not in self.earnings_surprise:
                            self.earnings_surprise[symbol] = []
                        
                        # surprise data is ready.
                        elif len(self.earnings_surprise[symbol]) >= self.min_surprise_period:
                            earnings_surprise_std: float = np.std(self.earnings_surprise[symbol])
                            sue: float = earnings_surprise / earnings_surprise_std
                            
                            # EAR calculation
                            min_day: datetime.date = (last_earnings_date - BDay(2)).date()
                            max_day: datetime.date = (last_earnings_date + BDay(1)).date()
                            stock_closes_around_earnings: List[Symbol] = [x for x in self.price_data_with_date[symbol] if x[0] >= min_day and x[0] <= max_day]
                            market_closes_around_earnings: List[Symbol] = [x for x in self.price_data_with_date[self.market] if x[0] >= min_day and x[0] <= max_day]
            
                            if len(stock_closes_around_earnings) == 4 and len(market_closes_around_earnings) == 4:
                                stock_return: float = stock_closes_around_earnings[-1][1] / stock_closes_around_earnings[0][1] - 1
                                market_return: float = stock_closes_around_earnings[-1][1] / stock_closes_around_earnings[0][1] - 1
                                
                                ear: float = stock_return - market_return
                                sue_ear[symbol] = (sue, ear)
        
                                # store pair in this month's history
                                self.sue_ear_history_actual.append((sue, ear))
    
                        self.earnings_surprise[symbol].append(earnings_surprise)
    
        # wait until we have history data for previous three months.
        if len(sue_ear) != 0 and len(self.sue_ear_history_previous) != 0:
            # Sort by SUE and EAR.
            sue_values: List[float] = [x[0] for x in self.sue_ear_history_previous]
            ear_values: List[float] = [x[1] for x in self.sue_ear_history_previous]
            
            top_sue_quantile: float = np.percentile(sue_values, self.percentile_range[0])
            bottom_sue_quantile: float = np.percentile(sue_values, self.percentile_range[1])
        
            top_ear_quantile: float = np.percentile(ear_values, self.percentile_range[0])
            bottom_ear_quantile: float = np.percentile(ear_values, self.percentile_range[1])
            
            self.long: List[Symbol] = [x[0] for x in sue_ear.items() if x[1][0] >= top_sue_quantile and x[1][1] >= top_ear_quantile]
        
        return self.long
        
    def OnData(self, data: Slice) -> None:
        # order execution
        targets: List[PortfolioTarget] = []
        for symbol in self.long:
            if symbol in data and data[symbol]:
                targets.append(PortfolioTarget(symbol, 1. / len(self.long)))
        
        self.SetHoldings(targets, True)

        self.long.clear()

    def Selection(self) -> None:
        self.selection_flag = True
        
        # store new EAR and SUE values every three months
        if self.month % 3 == 0:
            # Save previous month history.
            self.sue_ear_history_previous = self.sue_ear_history_actual
            self.sue_ear_history_actual.clear()

        self.month += 1
        if self.month > 12:
            self.month = 1

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