# region imports
from AlgorithmImports import *
from itertools import groupby
# endregion

class LiquidAssetsUniverseSelectionModel(FineFundamentalUniverseSelectionModel):
    def __init__(self, algorithm, universe_settings, universe_size):
        self.algorithm = algorithm
        self.universe_size = universe_size
        self.month = -1
        super().__init__(self._select_coarse, self._select_fine, universe_settings)

    def _select_coarse(self, coarse):
        # Monthly universe refresh
        if self.algorithm.time.month == self.month:
            return Universe.UNCHANGED
        self.month = self.algorithm.time.month
        
        selected = sorted([c for c in coarse if c.has_fundamental_data], 
                         key=lambda c: c.dollar_volume, reverse=True)[:self.universe_size]
        return [c.symbol for c in selected]

    def _select_fine(self, fine):
        return [f.symbol for f in fine]

class SectorRotationAlphaModel(AlphaModel):
    def __init__(self, num_sectors=3, stocks_per_sector=5):
        self.num_sectors = num_sectors
        self.stocks_per_sector = stocks_per_sector
        self.symbol_data = {}  # Stores symbol indicators and sector
        
    def update(self, algorithm, data):
        insights = []
        
        # Filter symbols with ready indicators
        ready_symbols = {symbol: sd for symbol, sd in self.symbol_data.items() 
                        if all([indicator.IsReady for indicator in sd['indicators'].values()])}
        
        if len(ready_symbols) == 0:
            return insights
        
        # Group by sector and calculate average momentum
        sectors = {}
        for symbol, sd in ready_symbols.items():
            sector = sd['sector']
            momentum = sd['indicators']['momentum'].Current.Value
            if sector not in sectors:
                sectors[sector] = []
            sectors[sector].append((symbol, momentum))
        
        # Calculate sector scores (average momentum)
        sector_scores = {sector: np.mean([m for _, m in symbols]) 
                        for sector, symbols in sectors.items()}
        top_sectors = sorted(sector_scores.items(), key=lambda x: x[1], reverse=True)[:self.num_sectors]
        
        # Select top stocks per sector
        for sector, _ in top_sectors:
            sector_stocks = sorted(sectors[sector], key=lambda x: x[1], reverse=True)[:self.stocks_per_sector]
            weight = 1 / (self.num_sectors * len(sector_stocks))
            for symbol, _ in sector_stocks:
                insights.append(Insight.Price(symbol, Expiry.END_OF_DAY, InsightDirection.UP, weight=weight))
        
        return insights

    def on_securities_changed(self, algorithm, changes):
        for security in changes.AddedSecurities:
            # Initialize indicators
            roc = algorithm.ROC(security.Symbol, 20, Resolution.Daily)
            self.symbol_data[security.Symbol] = {
                'indicators': {'momentum': roc},
                'sector': security.Fundamentals.AssetClassification.MorningstarSectorCode
            }
            
        for security in changes.RemovedSecurities:
            self.symbol_data.pop(security.Symbol, None)

class SectorRotation(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash(1000000)
        self.SetWarmUp(20, Resolution.Daily)  # Warm up for momentum indicator
        
        # Universe setup
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(LiquidAssetsUniverseSelectionModel(self, self.UniverseSettings, 100))
        
        # Alpha model (rotates top 3 sectors, 5 stocks each)
        self.AddAlpha(SectorRotationAlphaModel(num_sectors=3, stocks_per_sector=8))
        
        # Portfolio construction (daily rebalance)
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(Expiry.EndOfDay))
        
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())

    def OnData(self, data):
        pass  # Handled by alpha and portfolio models