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)
            # algorithm.Log(f"Added {added.symbol} to universe")

        for removed in changes.removed_securities:
            symbol = removed.symbol
            if symbol in self.symbols:
                self.symbols.remove(symbol)
                # algorithm.Log(f"Removed {symbol} from universe")

    def update(self, algorithm, data):
        # algorithm.Debug(f"Update method called for month {algorithm.time.month}, universe size: {len(self.symbols)}")
        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:

            # 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) + 10, Resolution.DAILY)
            # algorithm.Log(f"History size for {symbol}: {len(history)}")
            # 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] = max(momp.Current.Value / std.Current.Value, 0)
            # algorithm.Log(f"Processing symbol {symbol} with alpha value: {alphas[symbol]}")

        # Rank the symbol by the value of mom/vol
        selected = sorted(alphas.items(), key=lambda x: x[1], reverse=True)[:50]
        selected_symbols = [x[0] for x in selected]
        # algorithm.Debug(f"Selected symbols at {algorithm.Time}: {', '.join([str(symbol) for symbol in selected_symbols])}")        
        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(2020, 1, 1)
        self.set_end_date(2025, 1, 1)
        self.set_cash(1000000)            
        # Set number days to trace back
        v_lookback = self.get_parameter("v_lookback", 36)
        m_lookback = self.get_parameter("m_lookback", 12)
        
        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)

        # # DOW 30 Companies
        # dia = self.add_equity("DIA",
        #     resolution = self.universe_settings.resolution,
        #     data_normalization_mode = self.universe_settings.data_normalization_mode).symbol
        # self.set_benchmark(dia)
        
        self.set_universe_selection(etf_constituents_universe(spy, self.universe_settings))
        self.add_alpha(conservative_reblancing(spy, v_lookback, m_lookback))
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        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]