from AlgorithmImports import *

class conservative_reblancing(AlphaModel):

    def __init__(self, benchmark, v_lookback, m_lookback):
        self.benchmark = benchmark
        self.v_lookback = v_lookback
        self.m_lookback = m_lookback
        self.symbols = []
        self.month = -1

    def on_securities_changed(self, algorithm, changes):
        for added in changes.added_securities:
            self.symbols.append(added.symbol)

        for removed in changes.removed_securities:
            symbol = removed.symbol
            if symbol in self.symbols:
                self.symbols.remove(symbol)

    def update(self, algorithm, data):
        if algorithm.time.month == self.month: return []
        self.month = algorithm.time.month

        # Initialize the data
        alphas = dict()

        # Fetch indicator data
        for symbol in self.symbols:

            if symbol not in data.Keys:
                algorithm.Debug(f"At {algorithm.Time}, {symbol} is not in the data feed")
                continue
            

            # Create the indicators
            roc = algorithm.roc(symbol, 1, Resolution.Daily)
            std = algorithm.std(symbol, self.v_lookback, Resolution.DAILY)
            momp = algorithm.momp(symbol, self.m_lookback, Resolution.DAILY)

            # Get historical data for warm-up
            history = algorithm.History(symbol, max(self.v_lookback, self.m_lookback), Resolution.DAILY)

            # Warm up the indicators
            for idx, row in history.loc[symbol].iterrows():
                roc.Update(idx, row["close"])
                std.Update(idx, roc.current.value)
                momp.Update(idx, row["close"])

            # Compute the rank value
            alphas[symbol] = momp.Current.Value / std.Current.Value

        # Rank the symbol by the value of mom/vol
        selected = sorted(alphas.items(), key=lambda x: x[1], reverse=True)[:5]
        selected_symbols = [x[0] for x in selected]
        
        return [
            Insight.price(symbol, Expiry.END_OF_MONTH, InsightDirection.UP) for symbol in selected_symbols
        ]

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

class ConservativeApgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2018, 1, 1)
        self.set_end_date(2024, 1, 1)
        self.set_cash(100000)            
        # self.set_warm_up(60)
        # Set number days to trace back
        v_lookback = self.get_parameter("v_lookback", 36)
        m_lookback = self.get_parameter("m_lookback", 12)
        self.set_warmup(40)
        
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        # SPY 500 companies
        spy = self.add_equity("SPY",
            resolution = self.universe_settings.resolution,
            data_normalization_mode = self.universe_settings.data_normalization_mode).symbol
        self.set_benchmark(spy)
        
        self.set_universe_selection(etf_constituents_universe(spy, self.universe_settings))
        self.add_alpha(conservative_reblancing(spy, v_lookback, m_lookback))
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(Expiry.END_OF_MONTH))
        self.set_risk_management(NullRiskManagementModel())
        self.set_execution(ImmediateExecutionModel())

from AlgorithmImports import *

class etf_constituents_universe(ETFConstituentsUniverseSelectionModel):
    def __init__(self, benchmark, universe_settings: UniverseSettings = None) -> None:
        super().__init__(benchmark, universe_settings, self.etf_constituents_filter)

    def etf_constituents_filter(self, constituents):
        return [c.symbol for c in constituents]