# region imports
from AlgorithmImports import *
# endregion


class CustomHistoricalReturnsAlphaModel(AlphaModel):
    '''Uses Historical returns to create insights.'''

    def __init__(self, *args, **kwargs):
        '''Initializes a new default instance of the HistoricalReturnsAlphaModel class.
        Args:
            lookback(int): Historical return lookback period
            resolution: The resolution of historical data'''
        self.lookback = kwargs['lookback'] if 'lookback' in kwargs else 1
        self.resolution = kwargs['resolution'] if 'resolution' in kwargs else Resolution.DAILY
        self.prediction_interval = Time.multiply(Extensions.to_time_span(self.resolution), self.lookback)
        self._symbol_data_by_symbol = {}
        self.insight_collection = InsightCollection()

    def update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []

        for symbol, symbol_data in self._symbol_data_by_symbol.items():
            if symbol_data.can_emit:

                direction = InsightDirection.FLAT
                magnitude = symbol_data.return_
                if magnitude > 0: direction = InsightDirection.UP
                if magnitude < 0: direction = InsightDirection.DOWN

                if direction == InsightDirection.FLAT:
                    self.cancel_insights(algorithm, symbol)
                    continue

                insights.append(Insight.price(symbol, self.prediction_interval, direction, magnitude, weight=100*magnitude))

        self.insight_collection.add_range(insights)
        return insights

    def on_securities_changed(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''

        # clean up data for removed securities
        for removed in changes.removed_securities:
            symbol_data = self._symbol_data_by_symbol.pop(removed.symbol, None)
            if symbol_data is not None:
                symbol_data.remove_consolidators(algorithm)
            self.cancel_insights(algorithm, removed.symbol)

        # initialize data for added securities
        symbols = [ x.symbol for x in changes.added_securities ]
        history = algorithm.history(symbols, self.lookback, self.resolution)
        if history.empty: return

        tickers = history.index.levels[0]
        for ticker in tickers:
            symbol = SymbolCache.get_symbol(ticker)

            if symbol not in self._symbol_data_by_symbol:
                symbol_data = SymbolData(symbol, self.lookback)
                self._symbol_data_by_symbol[symbol] = symbol_data
                symbol_data.register_indicators(algorithm, self.resolution)
                symbol_data.warm_up_indicators(history.loc[ticker])

    def cancel_insights(self, algorithm, symbol):
        if not self.insight_collection.contains_key(symbol):
            return
        insights = self.insight_collection[symbol]
        algorithm.insights.cancel(insights)
        self.insight_collection.clear([ symbol ]);


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, symbol, lookback):
        self.symbol = symbol
        self.roc = RateOfChange('{}.roc({})'.format(symbol, lookback), lookback)
        self.consolidator = None
        self.previous = 0

    def register_indicators(self, algorithm, resolution):
        self.consolidator = algorithm.resolve_consolidator(self.symbol, resolution)
        algorithm.register_indicator(self.symbol, self.roc, self.consolidator)

    def remove_consolidators(self, algorithm):
        if self.consolidator is not None:
            algorithm.subscription_manager.remove_consolidator(self.symbol, self.consolidator)

    def warm_up_indicators(self, history):
        for tuple in history.itertuples():
            self.roc.update(tuple.Index, tuple.close)

    @property
    def return_(self):
        return float(self.roc.current.value)

    @property
    def can_emit(self):
        if self.previous == self.roc.samples:
            return False

        self.previous = self.roc.samples
        return self.roc.is_ready

    def __str__(self, **kwargs):
        return '{}: {:.2%}'.format(self.roc.name, (1 + self.return_)**252 - 1)

# region imports
from AlgorithmImports import *
from universe import PiotroskiScoreUniverseSelectionModel
from alpha import CustomHistoricalReturnsAlphaModel
# endregion


class PiotroskiScoreAlgorithm(QCAlgorithm):

    def initialize(self):
        # Algorithm cash and period setup
        self.set_cash(1_000_000_000)
        self.set_start_date(2016, 1, 1)
        
        # Configure settings to rebalance monthly.
        rebalance_date = self.date_rules.month_start("SPY")
        
        # Universe settings
        self.universe_settings.schedule.on(rebalance_date)
        self.universe_settings.resolution = Resolution.DAILY
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.add_universe_selection(PiotroskiScoreUniverseSelectionModel(
            self,
            self.get_parameter('score_threshold', 7),
            self.get_parameter('max_universe_size', 3)
        ))

        # Long only strategy on selected assets.
        self.add_alpha(CustomHistoricalReturnsAlphaModel(lookback=21))
        # Weight assets proportional to past returns.
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(rebalance_date, PortfolioBias.LONG))
        self.set_execution(ImmediateExecutionModel())
        self.set_warm_up(timedelta(3*365))
    

# region imports
from operator import attrgetter

from AlgorithmImports import *
# endregion


class PiotroskiFactors:

    field_map = {
        "roa": "operation_ratios.roa.one_year",
        "operating_cash_flow": "financial_statements.cash_flow_statement.cash_flow_from_continuing_operating_activities.twelve_months",
        "current_ratio": "operation_ratios.current_ratio.one_year",
        "ordinary_shares_number": "financial_statements.balance_sheet.ordinary_shares_number.twelve_months",
        "gross_margin": "operation_ratios.gross_margin.one_year",
        "assets_turnover": "operation_ratios.assets_turnover.one_year",
        "total_assets": "financial_statements.balance_sheet.total_assets.twelve_months",
        "long_term_debt": "financial_statements.balance_sheet.long_term_debt.twelve_months",
    }

    def __init__(self, f):
        for name, path in self.field_map.items():
            setattr(self, name, attrgetter(path)(f))

    @staticmethod
    def are_available(f):
        return all(
            not np.isnan(attrgetter(path)(f))
            for path in PiotroskiFactors.field_map.values()
        )

# region imports
from AlgorithmImports import *
# endregion


class PiotroskiScore:

    # Source: https://www.anderson.ucla.edu/documents/areas/prg/asam/2019/F-Score.pdf

    def get_score(self, factors):
        return (
            self.roa_score(factors)   # ROA
            + self.operating_cash_flow_score(factors)  # CFO
            + self.roa_change_score(factors)  # ΔROA
            + self.accruals_score(factors)  # ACCRUAL
            + self.leverage_score(factors)  # ΔLEVER
            + self.liquidity_score(factors)  # ΔLIQUID
            + self.share_issued_score(factors)  # EQ_OFFER
            + self.gross_margin_score(factors)  # ΔMARGIN
            + self.asset_turnover_score(factors)  # ΔTURN
        )

    def roa_score(self, factors):
        '''Get the Profitability - Return of Asset sub-score of Piotroski F-Score
        
        "I define ROA ... as net income before extraordinary items ..., scaled by 
        beginning of the year total assets." (p. 7)

        "Net income before extraordinary items for the fiscal year preceding 
        portfolio formation scaled by total assets at the beginning of year t."
        (p. 13)
        '''
        return int(factors[0].roa > 0)

    def operating_cash_flow_score(self, factors):
        '''Get the Profitability - Operating Cash Flow sub-score of Piotroski F-Score
        
        "I define ... CRO as ... cash flow from operations ..., scaled by 
        beginning of the year total assets." (p. 7)

        "Cash flow from operations scaled by total assets at the beginning 
        of year t". (p. 13)
        '''
        return int(factors[0].operating_cash_flow > 0)

    def roa_change_score(self, factors):
        '''Get the Profitability - Change in Return of Assets sub-score of Piotroski F-Score
        
        "I define ΔROA as the current year's ROA less the prior year's ROA. If ΔROA > 0,
        the indicator variable F_ΔROA equals one, zero otherwise." (p. 7)

        "Change in annual ROA for the year preceding portfolio formation. ΔROA is
        calculated as ROA for year t less the firm's ROA for year t-1." (p. 13)
        '''
        return int(factors[0].roa > factors[1].roa)

    def accruals_score(self, factors):
        '''Get the Profitability - Accruals sub-score of Piotroski F-Score
        
        "I define the variable ACCRUAL as current year's net income before 
        extraordinary items less cash flow from operations, scaled by 
        beginning of the year total assets. The indicator variable 
        F_ACCRUAL equals one if CFO > ROA, zero otherwise." (p. 7)

        "Net income before extraordinary items less cash flow from 
        operations, scaled by total assets at the beginning of year t." (p. 13)
        '''
        # Nearest Operating Cash Flow and ROA as current year data.
        operating_cashflow = factors[0].operating_cash_flow
        roa = factors[0].roa
        total_assets = factors[1].total_assets  # Beginning of year t.
        return int(operating_cashflow / total_assets > roa)

    def leverage_score(self, factors):
        '''Get the Leverage, Liquidity and Source of Funds - Change in Leverage sub-score of Piotroski F-Score
        
        "I measure ΔLEVER as the historical change in the ratio of total 
        long-term debt to average total assets, and view an increase (decrease)
        in financial leverage as a negative (positive) signal.... I define the
        indicator variable F_ΔLEVER to equal one (zero) if the firm's leverage 
        ratio fell (rose) in the year preceding portfolio formation." (p. 8)

        "Change in the firm's debt-to-assets ratio between the end of year t 
        and year t-1. The debt-to-asset ratio is defined as the firm's total
        long-term debt (including the portion of long-term debt classified 
        as current) scaled by average total assets." (p. 13)
        '''
        lt_t = factors[0].long_term_debt
        lt_t1 = factors[1].long_term_debt
        a_t = factors[0].total_assets
        a_t1 = factors[1].total_assets
        a_t2 = factors[2].total_assets

        avg_assets_t  = (a_t  + a_t1) / 2.0
        avg_assets_t1 = (a_t1 + a_t2) / 2.0

        leverage_t  = lt_t  / avg_assets_t
        leverage_t1 = lt_t1 / avg_assets_t1

        # Score 1 if leverage decreased
        return int(leverage_t < leverage_t1)

    def liquidity_score(self, factors):
        '''Get the Liquidity score
        
        "The variable ΔLIQUID measures the historical change in the firm's 
        current ratio between the current and prior year, where I define the 
        current ratio as the ratio of current assets to current liabilities 
        at fiscal year-end. I assume that an improvement in liquidity (i.e., 
        ΔLIQUID > 0) is a good signal about the firm's ability to service 
        current debt obligations. The indicator variable F_ΔLIQUID equals 
        one if the firm's liquidity improved, zero otherwise." (p. 8)

        "Change in the firm's current ratio between the end of year t and 
        year t-1. Current ratio is defined as total current assets divided 
        by total current liabilities." (p. 13)
        '''
        return int(factors[0].current_ratio > factors[1].current_ratio)

    def share_issued_score(self, factors):
        '''Get the share issued score
        
        "I define the indicator variable EQ_OFFER to equal one if the firm 
        did not issue common equity in the year preceding portfolio formation, 
        zero otherwise." (p. 8)
        '''
        return int(factors[0].ordinary_shares_number <= factors[1].ordinary_shares_number)

    def gross_margin_score(self, factors):
        '''Get the gross margin score
        
        "I define ΔMARGIN as the firm's current gross margin ratio (gross 
        margin scaled by total sales) less the prior year's gross margin 
        ratio.... The indicator variable F_ΔMARGIN equals one if ΔMARGIN 
        is positive, zero otherwise." (p. 8)

        "Gross margin (net sales less cost of good sold) for the year 
        preceding portfolio formation, scaled by net sales for the year, 
        less the firm's gross margin (scaled by net sales) from year t-1." 
        (p. 13)
        '''
        return int(factors[0].gross_margin > factors[1].gross_margin)

    def asset_turnover_score(self, factors):
        '''Get the asset turnover score
        
        "I define ΔTURN as the firm's current year asset turnover ratio 
        (total sales scaled by beginning of the year total assets) less
        the prior year's asset turnover ratio.... The indicator variable
        F_ΔTURN equals one if ΔTURN is positive, zero otherwise." (p. 8)

        "Change in the firm's asset turnover ratio between the end of 
        year t and year t-1. The asset turnover ratio is defined as net 
        sales scaled by average total assets for the year" (p. 13)
        '''
        return int(factors[0].assets_turnover > factors[1].assets_turnover)

# region imports
from AlgorithmImports import *
from piotroski_score import PiotroskiScore
from piotroski_factors import PiotroskiFactors
# endregion


class SymbolData:

    _piotroski_score = PiotroskiScore()

    def __init__(self, fundamental):
        self.is_ready = False
        self.score = None
        self._period_ending_date = datetime.min
        self._factors = RollingWindow(3)
        self.update(fundamental)
    
    def update(self, fundamental):
        period_ending_date = fundamental.financial_statements.period_ending_date.twelve_months
        # If a new 10K hasn't been released yet...
        if self._period_ending_date == period_ending_date:
            # If it's been 5 months since the previous fiscal year, rebalance this asset.
            if (self._factors.is_ready and 
                (fundamental.end_time - self._period_ending_date).days > 5*30):
                self.score = self._piotroski_score.get_score(self._factors)
                self.is_ready = True
        # When a new 10K is released, add the fundamental data to the RollingWindow.
        elif PiotroskiFactors.are_available(fundamental):
            # If we've missed a 10K, reset the RollingWindow.
            if period_ending_date.year - self._period_ending_date.year > 1:
                self._factors.reset()
                self.is_ready = False
            self._factors.add(PiotroskiFactors(fundamental))
            self._period_ending_date = period_ending_date
        return self.is_ready

# region imports
from AlgorithmImports import *
from symbol_data import SymbolData
from piotroski_factors import PiotroskiFactors
# endregion


class PiotroskiScoreUniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self, algorithm, threshold, max_universe_size=100):
        super().__init__(self._select_assets)
        self._algorithm = algorithm
        self._threshold = threshold
        self._max_universe_size = max_universe_size
        self._symbol_data_by_symbol = {}
        self._tickers = ['META', 'AMZN', 'AAPL', 'GOOG', 'NFLX']

    def _select_assets(self, fundamentals):
        fundamentals = [f for f in fundamentals if f.symbol.value in self._tickers]
        # Update the Piotroski factors of all assets.
        for f in fundamentals:
            self._symbol_data_by_symbol.setdefault(f.symbol, SymbolData(f)).update(f)

        # Wait until we have sufficient history.
        if self._algorithm.is_warming_up:
            return Universe.UNCHANGED

        # We only want stocks with fundamental data and price > $1.
        f_scores = {}
        for f in fundamentals:
            score = (self._symbol_data_by_symbol[f.symbol].score or 0)
            if score >= self._threshold and (f.market_cap or 0) > 1_000_000_000:
                f_scores[f.symbol] = score

        # Select stocks with the highest F-Score, and take the top x:
        top_symbols = [
            symbol for symbol, score in sorted(
                f_scores.items(), key=lambda x: x[1], reverse=True
            ) if score >= self._threshold
        ][:self._max_universe_size]

        self._algorithm.plot('F-Scores', 'Total', len(f_scores))
        self._algorithm.plot('F-Scores', 'Above Thresold', len(top_symbols))

        return top_symbols