# region imports
from AlgorithmImports import *
# endregion


class FundamentalFactorAlphaModel(AlphaModel):
    
    def __init__(self, algorithm, date_rule, num_assets, quality_weight, value_weight, size_weight):
        self._algorithm = algorithm
        self._num_assets = num_assets
        self._securities = []
        self._insights = []
        # Add a Scheduled Event to create new insights each month.
        algorithm.schedule.on(date_rule, algorithm.time_rules.at(8, 0), self._create_insights)
        # Normalize the quality, value, and size weights.
        weight_sum = sum([quality_weight, value_weight, size_weight])
        self._quality_weight = quality_weight / weight_sum
        self._value_weight = value_weight / weight_sum
        self._size_weight = size_weight / weight_sum

    # As assets enter and leave the universe, update our `_securities` list.
    def on_securities_changed(self, algorithm, changes):
        for security in changes.added_securities:
            self._securities.append(security)
        for security in changes.removed_securities:
            if security in self._securities:
                self._securities.remove(security)

    def _create_insights(self):
        # Assign quality, value, size score to each stock.
        quality_by_security = self._get_scores([
            (lambda x: x.fundamentals.operation_ratios.gross_margin.value, True, 2), 
            (lambda x: x.fundamentals.operation_ratios.quick_ratio.Value, True, 1), 
            (lambda x: x.fundamentals.operation_ratios.debt_to_assets.Value, False, 2)
        ])
        value_by_security = self._get_scores([
            (lambda x: x.fundamentals.valuation_ratios.book_value_per_share, True, 0.5),
            (lambda x: x.fundamentals.valuation_ratios.cash_return, True, 0.25),
            (lambda x: x.fundamentals.valuation_ratios.earning_yield, True, 0.25)
        ])
        size_by_security = self._get_scores([(lambda x: x.fundamentals.market_cap, False, 1)])
        # Assign a combined score to each stock.
        score_by_security = {
            symbol: (
                quality_value * self._quality_weight 
                + value_by_security[symbol] * self._value_weight 
                + size_by_security[symbol] * self._size_weight
            )
            for symbol, quality_value in quality_by_security.items()   
        }
        # Sort the securities by their scores.
        longs = sorted(score_by_security, key=lambda security: score_by_security[security])[:self._num_assets]
        # Create insights for each stock.
        self._insights = [
            Insight.price(security, Expiry.END_OF_MONTH, InsightDirection.UP, weight=(len(longs) - i)**2)
            for i, security in enumerate(longs)
        ]

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

    def _get_scores(self, fundamentals):
        '''Assigns scores to each stock in securities
        Args: 
            fundamentals: list of 3-tuples (lambda function, bool, float)
        Returns:
            Dictionary with score for each security'''
        # Normalize the factor weights.
        weights = [tup[2] for tup in fundamentals]
        weights = [float(i)/sum(weights) for i in weights]
        # Sort the securities by each factor.
        sorted_securities = [
            sorted(self._securities, key=fundamental, reverse=reverse)
            for fundamental, reverse, _ in fundamentals
        ]
        # Calculate a score for each stock.
        return {
            security: sum([
                sorted_securities[i].index(security) * weights[i] 
                for i in range(len(fundamentals))
            ])
            for security in self._securities    
        }
        

# region imports
from AlgorithmImports import *
# endregion
from AlphaModel import FundamentalFactorAlphaModel


class VerticalTachyonRegulators(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(1_000_000)
        self.settings.seed_initial_prices = True
        # Add universe selection.
        self._liquidity_filter_size = 200
        self.universe_settings.resolution = Resolution.DAILY
        date_rule = self.date_rules.month_start('SPY')
        self.universe_settings.schedule.on(date_rule)
        self.add_universe(self._select_assets)
        # Add an Alpha model.
        self.add_alpha(FundamentalFactorAlphaModel(self, date_rule, 20, 2, 2, 1))
        # Add a portfolio construction model.
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.settings.rebalance_portfolio_on_security_changes = False
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(date_rule))
        # Add a risk model.
        self.set_risk_management(TrailingStopRiskManagementModel(self.get_parameter("stopRisk", 0.1)))
        # Add an execution model.
        self.set_execution(ImmediateExecutionModel())
        # Plot the portfolio state each week.
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY), 
            self.time_rules.at(10, 30),
            lambda: self.plot(
                "Positions", 
                "Num", 
                len([symbol for symbol, holding in self.portfolio.items() if holding.invested])
            )
        )
        # Add a warm-up period so the algorithm trades on deployment.
        self.set_warm_up(timedelta(45))

    def _select_assets(self, fundamentals):        
        # Select only those with fundamental data and a sufficiently large price.
        filtered = [f for f in fundamentals if f.has_fundamental_data and f.price > 5]
        # Select the subset of stocks that are most liquid.
        filtered = sorted(filtered, key = lambda x: x.dollar_volume)[-self._liquidity_filter_size:]
        # Select the subset of stocks with the following fundamental factors:
        return [
            f.symbol for f in filtered 
            if (f.operation_ratios.gross_margin.value > 0 and 
                f.operation_ratios.quick_ratio.value > 0 and 
                f.operation_ratios.debt_to_assets.value > 0 and 
                f.valuation_ratios.book_value_per_share > 0 and 
                f.valuation_ratios.cash_return > 0 and 
                f.valuation_ratios.earning_yield > 0 and 
                f.market_cap > 0)
        ]