#region imports
from AlgorithmImports import *

from itertools import groupby
#endregion

class FactorSectorRotationAlphaModel(AlphaModel):
    
    securities = []
    new_factor_values = False
    month = -1

    def __init__(self, num_sectors: int = 2):
        self.num_sectors = num_sectors
        
    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Update factors
        for data_symbol, data_point in data.get(KavoutCompositeFactorBundle).items():
            self.new_factor_values = True
            underlying_security = algorithm.securities[data_symbol.underlying]
            # Create a composite factor from the first factor values
            underlying_security.value = data_point.growth + data_point.value_factor + data_point.quality + data_point.momentum + data_point.low_volatility
        if not self.new_factor_values:
            return []

        # Only emit insights when there is quote data, not when we get corporate action or alt data
        if data.quote_bars.count == 0:
            return []

        # Rebalance monthly
        if data.time.month == self.month:
            return []
        self.month = data.time.month
        self.new_factor_values = False

        # Calculate sector scores
        filtered = [security for security in self.securities if security.value is not None and security.symbol in data.quote_bars]
        key = lambda security: security.fundamentals.asset_classification.morningstar_sector_code
        sorted_by_industry_code = sorted(filtered, key=key)
        score_by_sector = {}
        for sector, securities in groupby(sorted_by_industry_code, key):
            values = [security.value for security in securities]
            score_by_sector[sector] = sum(values) / len(values)

        # Emit insights for the sectors with the highest aggregated scores
        insights = []
        selected_sectors = [k for k, _ in sorted(score_by_sector.items(), key = lambda x: x[1], reverse=True)[:self.num_sectors]]
        for sector, securities in groupby(sorted_by_industry_code, key):
            securities_copy = list(securities)
            # Split the portfolio evenly among the different sectors and give an equal-weight allocation to the constituents in each sector
            weight = 1 / len(securities_copy) / self.num_sectors
            for security in securities_copy:
                if sector in selected_sectors:
                    insights.append(Insight.price(security.symbol, Expiry.END_OF_MONTH, InsightDirection.UP, weight=weight))

        return insights


    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            # Subscribe to the Kavout Factor Bundle dataset
            security.dataset_symbol = algorithm.add_data(KavoutCompositeFactorBundle, security.symbol, Resolution.DAILY).symbol

            # Set the initial factor value to `None`
            security.value = None
            
            self.securities.append(security)

        # Remove the dataset subscription when the stock is removed from the universe
        for security in changes.removed_securities:
            if security in self.securities:
                algorithm.remove_security(security.dataset_symbol)
                self.securities.remove(security)