from AlgorithmImports import *


class FundamentalFactorAlphaModel(AlphaModel):
    
    def __init__(self, algorithm, date_rule):
        self._algorithm = algorithm
        self._equities_by_sector = {}
        # Update the insights at the start of each quarter.
        self._insights = []
        algorithm.schedule.on(date_rule, algorithm.time_rules.at(8, 0), self._create_insights)

    def _create_insights(self):
        # Update insights at the start of each quarter.
        if self._algorithm.time.month not in [1, 4, 7, 10]: 
            return
        self._insights.clear()
        # Iterate through each sector.
        for equities in self._equities_by_sector.values():
            # Rank the assets by their fundamental factors.
            sorted_by_roe = sorted(equities, key=lambda x: x.fundamentals.operation_ratios.roe.value)
            sorted_by_net_margin = sorted(equities, key=lambda x: x.fundamentals.operation_ratios.net_margin.value)
            sorted_by_pe = sorted(equities, key=lambda x: x.fundamentals.valuation_ratios.pe_ratio, reverse=True)
            # Calculate a score for each asset.
            score_by_equity = {
                equity: sum([
                    sorted_by_roe.index(equity), 
                    sorted_by_net_margin.index(equity), 
                    sorted_by_pe.index(equity)
                ])
                for equity in equities
            }
            # Buy the 20% of stocks (minimum 1) in this sector that had the greatest scores.
            length = max(int(len(score_by_equity)/5), 1)
            for security in sorted(score_by_equity, key=lambda e: score_by_equity[e])[-length:]:
                self._insights.append(Insight.price(security, Expiry.END_OF_QUARTER, InsightDirection.UP))

    def update(self, algorithm, data):
        insights = self._insights.copy()
        self._insights.clear()
        return insights

    def on_securities_changed(self, algorithm, changes):
        # As assets enter the universe, add them to `self._equities_by_sector`.
        for security in changes.added_securities:
            sector = security.fundamentals.asset_classification.morningstar_sector_code
            self._equities_by_sector.setdefault(sector, set()).add(security)
        # As assets leave the universe, remove them from `self._equities_by_sector`.
        for security in changes.removed_securities:
            sector = security.fundamentals.asset_classification.morningstar_sector_code
            if security in self._equities_by_sector[sector]:
                self._equities_by_sector[sector].remove(security)

from alpha import FundamentalFactorAlphaModel
from AlgorithmImports import *


class VerticalTachyonRegulators(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.settings.seed_initial_prices = True

        # Universe selection
        self._sectors = [
            MorningstarSectorCode.FINANCIAL_SERVICES,
            MorningstarSectorCode.REAL_ESTATE,
            MorningstarSectorCode.HEALTHCARE,
            MorningstarSectorCode.UTILITIES,
            MorningstarSectorCode.TECHNOLOGY
        ]
        date_rule = self.date_rules.month_start('SPY')
        self.universe_settings.schedule.on(date_rule)
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._select_assets)
        
        # Alpha Model
        self.add_alpha(FundamentalFactorAlphaModel(self, date_rule))

        # Portfolio construction model
        self.settings.rebalance_portfolio_on_security_changes = False
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(lambda time: None))
        
        # Risk model
        self.set_risk_management(NullRiskManagementModel())

        # Execution model
        self.set_execution(ImmediateExecutionModel())
        
        # Add a warm-up so that algorithm trades right away.
        self.set_warm_up(timedelta(120))

    def _select_assets(self, fundamentals):
        # Update the universe only at the start of each new quarter.
        if self.time.month not in [1, 4, 7, 10]: 
            return Universe.UNCHANGED
        # Select the most liquid stocks that have fundamental data and are trading above $5.
        selected = [f for f in fundamentals if f.has_fundamental_data and f.price > 5]
        selected = sorted(selected, key=lambda f: f.dollar_volume)[-500:]
        # Select the subset of stocks that IPO'd more than 5 years ago in the targeted sectors.
        return [
            f.symbol for f in selected 
            if (f.security_reference.ipo_date + timedelta(365*5) < self.time and
                f.asset_classification.MorningstarSectorCode in self._sectors and
                f.operation_ratios.roe.value > 0 and
                f.operation_ratios.net_margin.value > 0 and
                f.valuation_ratios.pe_ratio > 0)
        ]