# https://quantpedia.com/strategies/earnings-announcements-combined-with-stock-repurchases/
#
# The investment universe consists of stocks from NYSE/AMEX/Nasdaq (no ADRs, CEFs or REITs), bottom 25% of firms by market cap are dropped.
# Each quarter, the investor looks for companies that announce a stock repurchase program (with announced buyback for at least 5% of outstanding stocks) 
# during days -30 to -15 before the earnings announcement date for each company.
# Investor goes long stocks with announced buybacks during days -10 to +15 around an earnings announcement.
# The portfolio is equally weighted and rebalanced daily.
#
# QC Implementation changes:
#   - Universe consists of tickers, which have earnings annoucement.

#region imports
from AlgorithmImports import *
import numpy as np
from typing import List, Dict
from dataclasses import dataclass
from pandas.tseries.offsets import BDay
#endregion

class EarningsAnnouncementsCombinedWithStockRepurchases(QCAlgorithm):

    def Initialize(self) -> None:
        self.SetStartDate(2015, 1, 1)
        self.SetCash(100_000) 
        
        self.exchange_codes:List[str] = ['NYS', 'NAS', 'ASE']	

        self.selected: Dict[str, Symbol] = {}
        self.price: Dict[str, float] = {}
        self.managed_symbols: List[ManagedSymbol] = []
        self.earnings_universe: List[str] = []
        
        self.earnings: Dict[datetime.date, str] = {}
        self.buybacks: Dict[datetime.date, str] = {}
        
        self.max_traded_stocks: int = 40 # maximum number of trading stocks
        self.quantile: int = 4
        self.leverage: int = 5
        self.open_trade_offset: int = 10
        self.close_trade_offset: int = 15
        self.announcement_lookback: List[int] = [30, 15]
        self.earnings_last_date: Union[None, datetime.date] = None
        self.buybacks_last_date: Union[None, datetime.date] = None
        
        symbol: Symbol = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        # self.first_date: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()
            self.earnings_last_date = date
            self.earnings[date] = []
            
            # if not self.first_date: self.first_date = date

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

                self.earnings[date].append(ticker)
                
                if ticker not in self.earnings_universe:
                    self.earnings_universe.append(ticker)
        
        # load buyback dates
        csv_data: str = self.Download('data.quantpedia.com/backtesting_data/equity/BUY_BACKS.csv')
        lines: str = csv_data.split('\r\n')
        
        for line in lines[1:]: # skip header
            line_split: str = line.split(';')
            date: str = line_split[0]
            
            if date == '':
                continue
            
            date: datetime.date = datetime.strptime(date, "%d.%m.%Y").date()
            self.buybacks_last_date =  date
            self.buybacks[date] = []
            
            for ticker in line_split[1:]: # skip date in current line
                self.buybacks[date].append(ticker)
        
        self.months_counter: int = 0
        self.selection_flag: bool = False
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.FundamentalSelectionFunction)
        self.Schedule.On(self.DateRules.MonthStart(symbol), self.TimeRules.AfterMarketOpen(symbol), self.Selection)
        
    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(self.leverage)
            
    def FundamentalSelectionFunction(self, fundamental: List[Fundamental]) -> List[Symbol]:        
        # update stocks last prices
        for stock in fundamental:
            ticker: str = stock.Symbol.Value
            
            if ticker in self.earnings_universe:
                # store stock's last price
                self.price[ticker] = stock.AdjustedPrice
        
        # rebalance quarterly
        if not self.selection_flag:
            return Universe.Unchanged
        self.selection_flag = False
        
        # select stocks, which had spin off
        selected: List[Fundamental] = [x for x in fundamental if x.MarketCap != 0 \
                                and x.SecurityReference.ExchangeId in self.exchange_codes \
                                and x.Symbol.Value in self.earnings_universe]

        if len(selected) < self.quantile:
            return Universe.Unchanged
        
        # exclude 25% stocks with lowest market capitalization
        quantile: int = int(len(selected) / self.quantile)
        self.selected = {x.Symbol.Value : x.Symbol for x in sorted(selected, key=lambda x: x.MarketCap)[quantile:]}
        
        return list(self.selected.values())

    def OnData(self, data: Slice) -> None:
        remove_managed_symbols: List[ManagedSymbol] = []

        # check last date on custom data
        if any([self.Time.date() > date for date in [self.earnings_last_date, self.buybacks_last_date]]):
            self.Liquidate()
            return
                
        # maybe there should be BDay(15)
        # liquidate_date: datetime.date = self.Time.date() - timedelta(15)
        
        # check if bought stocks have 15 days after earnings annoucemnet
        for managed_symbol in self.managed_symbols:
            if (managed_symbol.earnings_date + BDay(self.close_trade_offset)).date() <= self.Time.date():
                remove_managed_symbols.append(managed_symbol)
                
                # liquidate stock by selling it's quantity
                self.MarketOrder(managed_symbol.symbol, -managed_symbol.quantity)
                
        # remove liquidated stocks from self.managed_symbols
        for managed_symbol in remove_managed_symbols:
            self.managed_symbols.remove(managed_symbol)
        
        # maybe there should be BDay(10)
        after_current: datetime.date = (self.Time + BDay(self.open_trade_offset)).date()
        
        if after_current in self.earnings:
            # this stocks has earnings annoucement after 10 days
            stocks_with_earnings: str = self.earnings[after_current]
            
            # 30 days before earnings annoucement
            buyback_start: datetime.date = (self.Time - BDay(self.announcement_lookback[0] - self.open_trade_offset)).date()
            # 15 days before earnings annoucement
            buyback_end: datetime.date = (self.Time - BDay(self.announcement_lookback[1] - self.open_trade_offset)).date()
            
            stocks_with_buyback: List[Symbol] = [] # storing stocks with buyback in period -30 to -15 days before earnings annoucement
            
            for buyback_date, tickers in self.buybacks.items():
                # check if buyback date is in period before earnings annoucement
                if buyback_date >= buyback_start and buyback_date <= buyback_end:
                    # iterate through each stock ticker for buyback date
                    for ticker in tickers:
                        # add stock ticker if it isn't already added, it has earnings annoucement after 10 days and was selected in selected
                        if (ticker not in stocks_with_buyback) and (ticker in stocks_with_earnings) and (ticker in self.selected):
                            stocks_with_buyback.append(self.selected[ticker])
                            
            # buying stocks buyback in period -30 to -15 days before earnings annoucement
            # and stocks, which have earnings date -10 days before current date
            for symbol in stocks_with_buyback:
                # check if there is a place in Portfolio for trading current stock
                if not len(self.managed_symbols) < self.max_traded_stocks:
                    continue
                
                # calculate stock quantity
                weight: float = self.Portfolio.TotalPortfolioValue / self.max_traded_stocks
                quantity: int = np.floor(weight / self.price[symbol.Value])
                
                if symbol in data and data[symbol]:
                    # go long stock
                    self.MarketOrder(symbol, quantity)
                
                    # store stock's ticker, earnings date and traded quantity
                    self.managed_symbols.append(ManagedSymbol(symbol, after_current, quantity))
        
    def Selection(self) -> None:
        # quarterly selection
        if self.months_counter % 3 == 0:
            self.selection_flag = True
        self.months_counter += 1
        
@dataclass
class ManagedSymbol():
    symbol: Symbol
    earnings_date: datetime.date
    quantity: int
        
# Custom fee model.
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters: OrderFeeParameters) -> OrderFee:
        fee: float = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))