#region imports
from AlgorithmImports import *
#endregion


class MonthlyLongAlphaModel(AlphaModel):
    
    _securities = []
    _month = -1

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Rebalance monthly when there is quote data, not when we get corporate action
        if data.time.month == self._month or data.quote_bars.count == 0:
            return []

        if not self._securities:
            algorithm.log('MonthlyLongAlphaModel.update: securities collection is empty')
            return []

        self._month = data.time.month

        # Emit insights
        return [Insight.price(security.symbol, Expiry.EndOfMonth, InsightDirection.UP) 
            for security in self._securities if security.symbol in data.quote_bars and security.price]

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.added_securities:            
            self._securities.append(security)

        for security in changes.removed_securities:
            if security in self._securities:
                self._securities.remove(security)

#region imports
from AlgorithmImports import *

from universe import TopStandardizedUnexpectedEarningsUniverseSelectionModel
from alpha import MonthlyLongAlphaModel
#endregion


class StandardizedUnexpectedEarningsAlgorithm(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. URL: https://www.jstor.org/stable/pdf/247321.pdf?casa_token=KHX3qwnGgTMAAAAA:
           ycgY-PzPfQ9uiYzVYeOF6yRDaNcRkL6IhLmRJuFpI6iWxsXJgB2BcM6ylmjy-g6xv-PYbXySJ_VxDpFETxw1PtKGUi1d91ce-h-V6CaL_SAAB56GZRQ
       [2] Hou, Xue and Zhang, 2018, Replicating Anomalies, Review of Financial Studies,
           URL: http://theinvestmentcapm.com/HouXueZhang2019RFS.pdf
    '''

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False

    def initialize(self):
        months_count = self.get_parameter("months_count", 36) # Number of months of rolling window object

        # Set backtest start date and warm-up period
        WARM_UP_FOR_LIVE_MODE = self.get_parameter("warm_up_for_live_mode", 1)
        MORNING_STAR_LIVE_MODE_HISTORY = timedelta(30) # US Fundamental Data by Morningstar is limited to the last 30 days
        if self.live_mode:
            self.set_warm_up(MORNING_STAR_LIVE_MODE_HISTORY)

        else: # Backtest mode
            if WARM_UP_FOR_LIVE_MODE: # Need to run a backtest before you can deploy this algorithm live
                now = datetime.now()
                self.set_start_date(now - MORNING_STAR_LIVE_MODE_HISTORY)
                self.set_end_date(now) # The universe selection model will quit 30-days before the current day
            else: # Regular backtest
                self.set_start_date(2023, 3, 1)
                self.set_end_date(2024, 3, 1)
            self.set_warm_up(timedelta(31 * (months_count+1)))
        self.set_cash(1_000_000)

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.schedule.on(self.date_rules.month_start())
        self.add_universe_selection(TopStandardizedUnexpectedEarningsUniverseSelectionModel(
            self,
            self.universe_settings,
            self.get_parameter("coarse_size", 50),        # Number of stocks to return from Coarse universe selection
            self.get_parameter("top_percent", 0.05),      # Percentage of symbols selected based on SUE sorting
            self.get_parameter("months_eps_change", 12),  # Number of months of lag to calculate eps change
            months_count,
            WARM_UP_FOR_LIVE_MODE
        ))

        self.add_alpha(MonthlyLongAlphaModel())

        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self._month = -1
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self._rebalance_func))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())        

    def _rebalance_func(self, time):
        if self._month != self.time.month and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self._month = self.time.month
            return time
        return None

    def on_data(self, data):
        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self._checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self._undesired_symbols_from_previous_deployment.append(symbol)
            self._checked_symbols_from_previous_deployment = True
        
        for symbol in self._undesired_symbols_from_previous_deployment[:]:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self._undesired_symbols_from_previous_deployment.remove(symbol)

#region imports
from AlgorithmImports import *
#endregion
# 27/10/2023: - Implement new Scheduled Universe function, and convert to pep8 syntax
#
# 27/10/2023: - Implement new Fundamental Universe Selection Model, merging coarse and fine selections
#
# 08/14/2023: -Reduced warm-up period to 30-days in live mode since the live feed for US fundamentals has a 30-day quota
#             -Reduced the `coarse_size` parameter to 50
#             -Adjusted the algorithm to utilize the ObjectStore to warm-up fundamental data older than 30-days
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_bc50228f65fd60a403e041ad494a58e2.html
#
# 04/15/2024: -Updated to PEP8 style
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_6b6be217580dab839425fc33c1bb181d.html

#region imports
from AlgorithmImports import *
#endregion


class TopStandardizedUnexpectedEarningsUniverseSelectionModel(FundamentalUniverseSelectionModel):
        
    _month = -1
    _hours = None
    _eps_by_symbol = {}   # Contains RollingWindow objects for all stocks
    _new_fine = []        # Contains new symbols selected at Coarse Selection
    
    def __init__(self, algorithm, universe_settings: UniverseSettings = None, coarse_size: int = 50, 
                top_percent: float = 0.05, months_eps_change: int = 12, months_count: int = 36, warm_up_for_live_mode: int = 1) -> None:
        def select(fundamental):
            # If it's a backtest and warm-up is over, save historical fundamental data into the ObjectStore so that you can use it in live mode
            if not algorithm.live_mode and not algorithm.is_warming_up and warm_up_for_live_mode:
                algorithm.object_store.save(self.OBJECT_STORE_MONTH_KEY, str(self._month))
                str_eps_by_symbol = {str(symbol): list(rolling_window)[::-1] for symbol, rolling_window in self._eps_by_symbol.items()}
                save_successful = algorithm.object_store.save(self.OBJECT_STORE_EPS_KEY, json.dumps(str_eps_by_symbol))
                algorithm.quit(f"Done warming up fundamentals. Save was successful: {save_successful}")
                return []
            
            if not self._hours or algorithm.live_mode:
                self._hours = algorithm.market_hours_database.get_entry(Market.USA, "SPY", SecurityType.EQUITY).exchange_hours
            
            selected = [ x for x in fundamental if x.has_fundamental_data and x.price > 5]
            sorted_by_dollar_volume = sorted(selected, key=lambda c: c.dollar_volume, reverse=True)
            self._new_fine = [c.symbol for c in sorted_by_dollar_volume[:coarse_size]]
            # Return all symbols that have appeared in Coarse Selection
            symbols = list( set(self._new_fine).union( set(self._eps_by_symbol.keys()) ) )

            sue_by_symbol = dict()
            
            for stock in fundamental:
                if stock.symbol not in symbols:
                    continue

                ### Save (symbol, rolling window of EPS) pair in dictionary
                if not stock.symbol in self._eps_by_symbol:
                    self._eps_by_symbol[stock.symbol] = RollingWindow[float](months_count)
                # update rolling window for each stock
                self._eps_by_symbol[stock.symbol].add(stock.earning_reports.basic_eps.three_months)
            
                ### Calculate SUE
            
                if stock.symbol in self._new_fine and self._eps_by_symbol[stock.symbol].is_ready:
            
                    # Calculate the EPS change from four quarters ago
                    rw = self._eps_by_symbol[stock.symbol]
                    eps_change = rw[0] - rw[months_eps_change]
                    
                    # Calculate the st dev of EPS change for the prior eight quarters
                    eps_list = list(rw)[::-1]
                    new_eps_list = eps_list[:months_count - months_eps_change:3]
                    old_eps_list = eps_list[months_eps_change::3]
                    eps_std = np.std( [ new_eps - old_eps for new_eps, old_eps in 
                                        zip( new_eps_list, old_eps_list )
                                    ] )
                    
                    # Get Standardized Unexpected Earnings (SUE)
                    sue_by_symbol[stock.symbol] = eps_change / eps_std
            
            # Sort and return the top quantile
            sorted_dict = sorted(sue_by_symbol.items(), key = lambda x: x[1], reverse = True)
            symbols = [ x[0] for x in sorted_dict[:math.ceil( top_percent * len(sorted_dict) )] ]
            return symbols

        super().__init__(select, universe_settings)

        # If it's live mode, load the historical fundamental data from the ObjectStore
        self.OBJECT_STORE_EPS_KEY = f"{algorithm.project_id}/fundamentals-warm-up"
        self.OBJECT_STORE_MONTH_KEY = f"{algorithm.project_id}/month"
        if algorithm.live_mode:
            if not algorithm.object_store.contains_key(self.OBJECT_STORE_EPS_KEY):
                algorithm.quit("No fundamental data in the ObjectStore. Run a backtest before deploying live.")
                return
            self._eps_by_symbol = {}
            for security_id, eps_list in json.loads(algorithm.object_store.read(self.OBJECT_STORE_EPS_KEY)).items():
                window = RollingWindow[float](months_count)
                for x in eps_list:
                    window.add(x)
                self._eps_by_symbol[algorithm.symbol(security_id)] = window
            self._month = int(algorithm.object_store.read(self.OBJECT_STORE_MONTH_KEY))