#region imports
from AlgorithmImports import *
#endregion


class MeanReversionAlphaModel(AlphaModel):

    _securities = []
    _month = -1

    def __init__(self, roc_period, num_positions_per_side):
        self._roc_period = roc_period
        self._num_positions_per_side = num_positions_per_side

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Reset indicators when corporate actions occur
        for symbol in set(data.splits.keys() + data.dividends.keys()):
            security = algorithm.securities[symbol]
            if security in self._securities:
                algorithm.unregister_indicator(security.indicator)
                self._initialize_indicator(algorithm, security)
        
        # Only emit insights when there is quote data, not when a corporate action occurs (at midnight)
        if data.quote_bars.count == 0:
            return []

        # Only emit insights once per month
        if self._month == algorithm.time.month:
            return []
        
        # Check if enough indicators are ready
        ready_securities = [security for security in self._securities if security.indicator.is_ready and security.symbol in data.quote_bars]
        if len(ready_securities) < 2 * self._num_positions_per_side:
            return []

        self._month = algorithm.time.month

        # Short securities that have the highest trailing ROC
        sorted_by_roc = sorted(ready_securities, key=lambda security: security.indicator.current.value)
        insights = [Insight.price(security.symbol, Expiry.END_OF_MONTH, InsightDirection.DOWN) for security in sorted_by_roc[-self._num_positions_per_side:]]
        # Long securities that have the lowest trailing ROC
        insights += [Insight.price(security.symbol, Expiry.END_OF_MONTH, InsightDirection.UP) for security in sorted_by_roc[:self._num_positions_per_side]]
        return insights

    def _initialize_indicator(self, algorithm, security):
        security.indicator = algorithm.ROC(security.symbol, self._roc_period, Resolution.DAILY)
        algorithm.warm_up_indicator(security.symbol, security.indicator)

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

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

#region imports
from AlgorithmImports import *

from universe import CountryEquityIndexUniverseSelectionModel
from alpha import MeanReversionAlphaModel
#endregion


class CountryEquityIndexesMeanReversionAlgorithm(QCAlgorithm):

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False
    _previous_expiry_time = None

    def initialize(self):
        self.set_start_date(2023, 3, 1)  # Set Start Date
        self.set_end_date(2024, 3, 1) 
        self.set_cash(1_000_000) 

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.add_universe_selection(CountryEquityIndexUniverseSelectionModel())
        
        self.add_alpha(MeanReversionAlphaModel(
            self.get_parameter("roc_period_months", 6) * 21,
            self.get_parameter("num_positions_per_side", 5)
        ))
        
        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self._rebalance_func))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        self.set_warm_up(timedelta(31))

    def _rebalance_func(self, time):
        # Rebalance when all of the following are true:
        # - There are new insights or old insights have been cancelled since the last rebalance
        # - The algorithm isn't warming up
        # - There is QuoteBar data in the current slice
        latest_expiry_time = sorted([insight.close_time_utc for insight in self.insights], reverse=True)[0] if self.insights.count else None
        if self._previous_expiry_time != latest_expiry_time and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self._previous_expiry_time = latest_expiry_time
            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


class CountryEquityIndexUniverseSelectionModel(ManualUniverseSelectionModel):
    def __init__(self):
        tickers = [
            "EWJ",  # iShares MSCI Japan Index ETF
            "EFNL", # iShares MSCI Finland Capped Investable Market Index ETF
            "EWW",  # iShares MSCI Mexico Inv. Mt. Idx
            "ERUS", # iShares MSCI Russia ETF
            "IVV",  # iShares S&P 500 Index
            "AUD",  # Australia Bond Index Fund
            "EWQ",  # iShares MSCI France Index ETF
            "EWH",  # iShares MSCI Hong Kong Index ETF
            "EWI",  # iShares MSCI Italy Index ETF
            "EWY",  # iShares MSCI South Korea Index ETF
            "EWP",  # iShares MSCI Spain Index ETF
            "EWD",  # iShares MSCI Sweden Index ETF
            "EWL",  # iShares MSCI Switzerland Index ETF
            "EWC",  # iShares MSCI Canada Index ETF
            "EWZ",  # iShares MSCI Brazil Index ETF
            "EWO",  # iShares MSCI Austria Investable Mkt Index ETF
            "EWK",  # iShares MSCI Belgium Investable Market Index ETF
            "BRAQ", # Global X Brazil Consumer ETF
            "ECH"   # iShares MSCI Chile Investable Market Index ETF
        ]
        symbols = [Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers]
        super().__init__(symbols)