#region imports
from AlgorithmImports import *

from collections import deque
import calendar
#endregion


class BetaAlphaModel(AlphaModel):

    def __init__(self, algorithm, benchmark, date_rule):
        self._benchmark = benchmark
        self._securities = []
        self._insights = []
        # Add a Scheduled Event to emit monthly insights.
        algorithm.schedule.on(date_rule, algorithm.time_rules.at(8, 0), self._create_insights)

    def on_securities_changed(self, algorithm, changes):
        # As assets enter the universe, create their beta indicators.
        for security in changes.added_securities:
            if security != self._benchmark:
                security.beta = algorithm.b(security, self._benchmark, 252)
                self._securities.append(security)
        # As assets leave the universe, stop updating their beta indicators.
        for security in changes.removed_securities:
            algorithm.deregister_indicator(security.beta)
            self._securities.remove(security)
    
    def _create_insights(self):
        # Sort the securities by their beta factor.
        sorted_by_beta = sorted(
            [security for security in self._securities if security.beta.is_ready], 
            key=lambda security: security.beta
        )
        # Create insights that will form a long-short portfolio.
        positions_per_side = int(0.25*len(sorted_by_beta))
        weight = 0.5 / positions_per_side
        self._insights = [
            Insight.price(security, Expiry.END_OF_MONTH, InsightDirection.UP, weight=weight)
            for security in sorted_by_beta[:positions_per_side]
        ] + [
            Insight.price(security, Expiry.END_OF_MONTH, InsightDirection.DOWN, weight=weight)
            for security in sorted_by_beta[-positions_per_side:]
        ]

    def update(self, algorithm, data):
        # Return our monthly insights for the portfolio construction model.
        insights = self._insights.copy()
        self._insights = []
        return insights

#region imports
from AlgorithmImports import *
from alpha import BetaAlphaModel
#endregion


class BetaAlgorithm(QCAlgorithmFramework):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        # Add a universe of country index ETFs.
        self.universe_settings.resolution = Resolution.DAILY
        tickers =  [
            "EWJ",  # iShares MSCI Japan Index ETF
            "EZU",  # iShares MSCI Eurozone 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
            "ICOL", # Consumer Discretionary Select Sector SPDR Fund
            "AAXJ", # iShares MSCI All Country Asia ex Japan Index ETF
            "AUD",  # Australia Bond Index Fund
            "EWQ",  # iShares MSCI France Index ETF
            "BUND", # Pimco Germany Bond Index Fund
            "EWH",  # iShares MSCI Hong Kong Index ETF
            "EPI",  # WisdomTree India Earnings ETF
            "EIDO"  # iShares MSCI Indonesia Investable Market Index ETF
            "EWI",  # iShares MSCI Italy Index ETF
            "GAF",  # SPDR S&P Emerging Middle East & Africa ETF
            "ENZL", # iShares MSCI New Zealand Investable Market Index Fund
            "NORW"  # Global X FTSE Norway 30 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
            "GXC",  # SPDR S&P China ETF
            "EWC",  # iShares MSCI Canada Index ETF
            "EWZ",  # iShares MSCI Brazil Index ETF
            "ARGT", # Global X FTSE Argentina 20 ETF
            "AND",  # Global X FTSE Andean 40 ETF
            "AIA",  # iShares S&P Asia 50 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
            "CHIB", # Global X China Technology ETF
            "EGPT", # Market Vectors Egypt Index ETF
            "ADRU"  # BLDRS Europe 100 ADR Index ETF
        ] 
        symbols = [Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers]
        self.set_universe_selection(ManualUniverseSelectionModel(symbols))
        # Add an Alpha model that emits insights based on beta.
        self.settings.automatic_indicator_warm_up = True
        benchmark = self.add_equity("SPY", Resolution.DAILY)
        date_rule = self.date_rules.month_start(benchmark)
        self.set_alpha(BetaAlphaModel(self, benchmark, date_rule))
        # Add a portfolio construction model.
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.settings.rebalance_portfolio_on_security_changes = False
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(date_rule))
        # Add the risk and execution models.
        self.set_risk_management(NullRiskManagementModel())
        self.set_execution(ImmediateExecutionModel())
        # Add a warm-up period so the algorithm trades right away.
        self.set_warm_up(timedelta(45))