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

class SymbolData():
    def __init__(self, period:int) -> None:
        self.prices:RollingWindow = RollingWindow[float](period)
        
    def update(self, price:float) -> None:
        self.prices.Add(price)
        
    def is_ready(self) -> bool:
        return self.prices.IsReady
        
    def performance(self) -> float:
        prices:list[float] = list(self.prices)
        return (prices[0] - prices[-1]) / prices[-1]

class ManagedSymbol():
    def __init__(self, symbol:Symbol, date_to_switch:datetime.date, date_to_liquidate:datetime.date) -> None:
        self.symbol:Symbol = symbol
        self.date_to_switch:datetime.date = date_to_switch
        self.date_to_liquidate:datetime.date = date_to_liquidate

class QuantpediaEarningsEps(PythonData):
    _earnings_universe:Set[str] = set()

    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource('data.quantpedia.com/backtesting_data/economic/{0}.json'.format(config.Symbol.Value.lower()), SubscriptionTransportMedium.RemoteFile, FileFormat.UnfoldingCollection)    
   
    @staticmethod
    def get_earnings_universe() -> list:
       return list(QuantpediaEarningsEps._earnings_universe)

    def Reader(self, config, line, date, isLiveMode):
        objects:list[QuantpediaEarningsEps] = []
        data:list[dict] = json.loads(line)
        end_time:datetime.date|None = None

        for index, sample in enumerate(data):
            custom_data:QuantpediaEarningsEps = QuantpediaEarningsEps()
            custom_data.Symbol = config.Symbol
            
            earnings_date:datetime.date = datetime.strptime(sample['date'], '%Y-%m-%d')
            # strategy trades 5 days before earnings day
            before_earnings_date:datetime.date = (earnings_date - BDay(5)).date()

            custom_data['earnings_date'] = earnings_date
            custom_data.Time = before_earnings_date
            custom_data.EndTime = custom_data.Time + timedelta(days=1)
            end_time = custom_data.EndTime
            
            curr_stocks:dict[str, dict] = {}

            for stock_data in sample['stocks']:
                ticker:str = stock_data['ticker']
                QuantpediaEarningsEps._earnings_universe.add(ticker)
                curr_stocks[ticker] = { attribute: value for attribute, value in stock_data.items() }

            custom_data['curr_earnings_stocks'] = curr_stocks
            objects.append(custom_data)

        return BaseDataCollection(end_time, config.Symbol, objects)

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

# https://quantpedia.com/strategies/post-earnings-announcement-drift-combined-with-strong-momentum/
#
# The investment universe consists of all stocks from NYSE, AMEX and NASDAQ with a price greater than $5. Each quarter, all stocks are
# sorted into deciles based on their 12 months past performance. The investor then uses only stocks from the top momentum decile and 
# goes long on each stock 5 days before the earnings announcement and closes the long position at the close of the announcement day. 
# Subsequently, at the close of the announcement day, he/she goes short and he/she closes his short position on the 5th day after the
# earnings announcement.
#
# QC Implementation changes:
#   - Investment universe consist of stocks with earnings data available.

from pandas.tseries.offsets import BDay
from AlgorithmImports import *
import data_tools

class PostEarningsAnnouncementDriftCombinedwithStrongMomentum(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2010, 1, 1)   # earnings days data starts in 2010
        self.SetCash(100_000)
        
        self.quantile: int = 10
        self.min_share_price: int = 5

        self.period: int = 12 * 21      # need n daily prices
        self.rebalance_period: int = 3  # referes to months, which has to pass, before next portfolio rebalance

        self.leverage: int = 5

        self.data: Dict[Symbol, data_tools.SymbolData] = {} 
        self.selected_symbols: List[Symbol] = []
        
        # 50 equally weighted brackets for traded symbols
        self.managed_symbols_size: int = 50
        self.managed_symbols: List[data_tools.ManagedSymbol] = []

        # earning data parsing
        self.earnings: Dict[datetime.date, list[str]] = {}
        days_before_earnings: List[datetime.date] = []
        
        earnings_set: Set(str) = set()
        
        # Source: https://www.nasdaq.com/market-activity/earnings
        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[date] = []
            days_before_earnings.append((date - BDay(5)).date())
            
            for stock_data in obj['stocks']:
                ticker: str = stock_data['ticker']

                self.earnings[date].append(ticker)
                earnings_set.add(ticker)

        self.earnings_universe: List[str] = list(earnings_set)
        
        self.symbol: Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol

        self.months_counter: int = 0
        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)
        
        # Events on earnings days, before and after earning days.
        self.Schedule.On(self.DateRules.On(days_before_earnings), self.TimeRules.AfterMarketOpen(self.symbol), self.DaysBefore)
        self.Schedule.On(self.DateRules.MonthStart(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol), self.Selection)
    
    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetFeeModel(data_tools.CustomFeeModel())
            security.SetLeverage(self.leverage)
                
    def FundamentalSelectionFunction(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # daily update of prices
        for stock in fundamental:
            symbol: Symbol = stock.Symbol

            if symbol in self.data:
                self.data[symbol].update(stock.AdjustedPrice)
        
        if not self.selection_flag:
            return Universe.Unchanged
        self.selection_flag = False

        selected: List[Symbol] = [
            x.Symbol for x in fundamental if x.HasFundamentalData 
            and x.Market == 'usa' and x.Price > self.min_share_price
            and x.Symbol.Value in self.earnings_universe
        ]
                                
        # warm up prices
        for symbol in selected:
            if symbol in self.data:
                continue
        
            self.data[symbol] = data_tools.SymbolData(self.period)
            history: DataFrame = self.History(symbol, self.period, Resolution.Daily)

            if history.empty:
                self.Log(f"Not enough data for {symbol} yet")
                continue

            closes: Series = history.loc[symbol].close
            for _, close in closes.items():
                self.data[symbol].update(close)
        
        # calculate momentum for each stock in self.earnings_universe
        momentum: Dict[Symbol, float] = {
            symbol: self.data[symbol].performance() for symbol in selected if self.data[symbol].is_ready()
        }

        if len(momentum) < self.quantile:
            self.selected_symbols = []
            return Universe.Unchanged
            
        quantile: int = int(len(momentum) / self.quantile)
        sorted_by_mom: List[Symbol] = sorted(momentum, key=momentum.get)

        # the investor uses only stocks from the top momentum quantile
        self.selected_symbols = sorted_by_mom[-quantile:]
        
        return self.selected_symbols

    def DaysBefore(self) -> None:
        # every day check if 5 days from now is any earnings day
        earnings_date: datetime.date = (self.Time + BDay(5)).date()
        date_to_liquidate: datetime.date = (earnings_date + BDay(6)).date()
        
        if earnings_date not in self.earnings:
            return

        for symbol in self.selected_symbols:
            ticker: str = symbol.Value
            # is there any symbol which has earnings in 5 days
            if ticker not in self.earnings[earnings_date]:
                continue

            if (len(self.managed_symbols) < self.managed_symbols_size) and not self.Securities[symbol].Invested and \
                self.Securities[symbol].Price != 0 and self.Securities[symbol].IsTradable:
                self.SetHoldings(symbol, 1 / self.managed_symbols_size)
                
                # NOTE: Must offset date to switch position by one day due to midnight execution of OnData function.
                # Alternatively, there's is a possibility to switch to BeforeMarketClose function.
                self.managed_symbols.append(data_tools.ManagedSymbol(symbol, (earnings_date + BDay(1)).date(), date_to_liquidate))
                    
    def OnData(self, data: Slice) -> None:
        # switch positions on earnings days.
        curr_date: datetime.date = self.Time.date()
        
        managed_symbols_to_delete: List[data_tools.ManagedSymbol] = []
        for managed_symbol in self.managed_symbols:
            if managed_symbol.date_to_switch == curr_date:
                # switch position from long to short
                if managed_symbol.symbol in data and data[managed_symbol.symbol]:
                    self.SetHoldings(managed_symbol.symbol, -1 / self.managed_symbols_size)
            
            elif managed_symbol.date_to_liquidate <= curr_date:
                self.Liquidate(managed_symbol.symbol)
                managed_symbols_to_delete.append(managed_symbol)
                
        # remove symbols from management
        for managed_symbol in managed_symbols_to_delete:
            self.managed_symbols.remove(managed_symbol)
            
    def Selection(self) -> None:
        # quarter selection
        if self.months_counter % self.rebalance_period == 0:
            self.selection_flag = True
        self.months_counter += 1