#region imports
from AlgorithmImports import *
#endregion


class LongMonthlyAlphaModel(AlphaModel):

    _securities = []

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Only generate long insights in January
        if algorithm.time.month == 1:
            return [Insight.price(security.symbol, Expiry.END_OF_MONTH, InsightDirection.UP) for security in self._securities]
        return []

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.removed_securities:
            if security in self._securities:
                self._securities.remove(security)
        self._securities.extend(changes.added_securities)

# January Effect in Stocks Strategy
# Exploits the well-documented January Effect anomaly where small-cap stocks
# tend to outperform large-caps in January due to tax-loss harvesting reversals,
# window dressing effects, and renewed investor optimism at the start of the year.
# The algorithm uses a custom universe selection model (JanuaryEffectUniverseSelection)
# to filter a coarse universe of 1,000 stocks down to the top 10 small-cap candidates
# based on fundamental data. Positions are entered at the start of January and held
# through the month, then liquidated. The strategy uses equal-weight allocation
# across selected securities with insight-driven portfolio construction.

#region imports
from AlgorithmImports import *

from universe import JanuaryEffectUniverseSelectionModel
from alpha import LongMonthlyAlphaModel
#endregion


class JanuaryEffectInStocksAlgorithm(QCAlgorithm):

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False

    def initialize(self):
        self.set_start_date(1998, 1, 1)  # Earliest available for US Equities fundamental data
        self.set_cash(100000) 

        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.schedule.on(self.date_rules.month_start())
        self.add_universe_selection(JanuaryEffectUniverseSelectionModel(
            self,
            self.universe_settings,
            self.get_parameter("coarse_size", 1_000),
            self.get_parameter("fine_size", 10)
        ))

        self.add_alpha(LongMonthlyAlphaModel())

        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.month = -1
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self._rebalance_func))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        self.set_warm_up(timedelta(31))

    def _rebalance_func(self, time):
        if self.month != self.time.month and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self.month = self.time.month
            return time
        return None
    
    def on_data(self, data):
        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self._checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self._undesired_symbols_from_previous_deployment.append(symbol)
            self._checked_symbols_from_previous_deployment = True
        
        for symbol in self._undesired_symbols_from_previous_deployment:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self._undesired_symbols_from_previous_deployment.remove(symbol)

#region imports
from AlgorithmImports import *
#endregion


class JanuaryEffectUniverseSelectionModel(FundamentalUniverseSelectionModel):
    def __init__(self, algorithm: QCAlgorithm, universe_settings: UniverseSettings = None, coarse_size: int = 1_000, fine_size: int = 10) -> None:
        def select(fundamental):
            # Select the securities that have the most dollar volume
            shortlisted = [c for c in sorted(fundamental, key=lambda x: x.dollar_volume, reverse=True)[:coarse_size]]
            
            fine = [i for i in shortlisted if i.earning_reports.basic_average_shares.three_months!=0
                                    and i.earning_reports.basic_eps.twelve_months!=0
                                    and i.valuation_ratios.pe_ratio!=0]
            # Sort securities by market cap
            sorted_by_market_cap = sorted(fine, key = lambda x: x.market_cap, reverse=True)
            # In January, select the securities with the smallest market caps
            if algorithm.time.month == 1:
                return [f.symbol for f in sorted_by_market_cap[-fine_size:]]
            # If it's not January, select the securities with the largest market caps
            return [f.symbol for f in sorted_by_market_cap[:fine_size]]
        
        super().__init__(select, universe_settings)