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


class PiotroskiScoreAlgorithm(QCAlgorithm):

    def initialize(self):
        # Algorithm cash and period setup
        self.set_cash(10_000_000)
        self.set_start_date(2022, 10, 1)
        self.set_end_date(2025, 10, 1)
        
        # Configure settings to rebalance monthly.
        rebalance_date = self.date_rules.month_start(Symbol.create('SPY', SecurityType.EQUITY, Market.USA))
        
        # Universe settings
        self.universe_settings.schedule.on(rebalance_date)
        self.universe_settings.resolution = Resolution.HOUR
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.add_universe_selection(PiotroskiScoreUniverseSelectionModel(
            self.get_parameter("score_threshold", 7),
            self.get_parameter("roa_multiple", 1),
            self.get_parameter('universe_size', 100)
        ))

        # Long only strategy on selected assets
        self.add_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(30)))
        # Weight sectors equally
        self.set_portfolio_construction(SectorWeightingPortfolioConstructionModel(rebalance_date))
        # Avoid illiquid assets. Maximum 1% spread allowed before execution 
        self.set_execution(SpreadExecutionModel(0.01))

# region imports
from AlgorithmImports import *
# endregion


class PiotroskiScore:

    def __init__(self, roa_multiple):
        self._roa_multiple = roa_multiple

    def get_score(self, fundamental):
        return (
            self.get_r_o_a_score(fundamental)
            + self.get_operating_cash_flow_score(fundamental)
            + self.get_r_o_a_change_score(fundamental)
            + self.get_accruals_score(fundamental)
            + self.get_leverage_score(fundamental)
            + self.get_liquidity_score(fundamental)
            + self.get_share_issued_score(fundamental)
            + self.get_gross_margin_score(fundamental)
            + self.get_asset_turnover_score(fundamental)
        )

    def get_r_o_a_score(self, fundamental, classify=True):
        '''Get the Profitability - Return of Asset sub-score of Piotroski F-Score'''
        # Nearest ROA as current year data
        roa = fundamental.operation_ratios.ROA.three_months
        if not classify:
            return roa
        # 1 score if ROA datum exists and positive, else 0
        return self._roa_multiple * int(roa and roa > 0)

    def get_operating_cash_flow_score(self, fundamental, classify=True):
        '''Get the Profitability - Operating Cash Flow sub-score of Piotroski F-Score'''
        # Nearest Operating Cash Flow as current year data
        operating_cashflow = fundamental.financial_statements.cash_flow_statement.cash_flow_from_continuing_operating_activities.three_months
        if not classify:
            return operating_cashflow
        # 1 score if operating cash flow datum exists and positive, else 0
        return int(operating_cashflow and operating_cashflow > 0)

    def get_r_o_a_change_score(self, fundamental, classify=True):
        '''Get the Profitability - Change in Return of Assets sub-score of Piotroski F-Score'''
        # if current or previous year's ROA data does not exist, return 0 score
        roa = fundamental.operation_ratios.ROA
        if not roa.three_months or not roa.one_year:
            return 0
        if not classify:
            return roa.three_months - roa.one_year
        # 1 score if change in ROA positive, else 0 score
        return int(roa.three_months > roa.one_year)

    def get_accruals_score(self, fundamental, classify=True):
        '''Get the Profitability - Accruals sub-score of Piotroski F-Score'''
        # Nearest Operating Cash Flow, Total Assets, ROA as current year data
        operating_cashflow = fundamental.financial_statements.cash_flow_statement.cash_flow_from_continuing_operating_activities.three_months
        total_assets = fundamental.financial_statements.balance_sheet.total_assets.three_months
        roa = fundamental.operation_ratios.ROA.three_months
        if not classify:
            return operating_cashflow / total_assets - roa if total_assets else None
        # 1 score if operating cash flow, total assets and ROA exists, and operating cash flow / total assets > ROA, else 0
        return int(operating_cashflow and total_assets and roa and operating_cashflow / total_assets > roa)

    def get_leverage_score(self, fundamental, classify=True):
        '''Get the Leverage, Liquidity and Source of Funds - Change in Leverage sub-score of Piotroski F-Score'''
        # if current or previous year's long term debt to equity ratio data does not exist, return 0 score
        long_term_debt_ratio = fundamental.operation_ratios.long_term_debt_equity_ratio
        if not long_term_debt_ratio.three_months or not long_term_debt_ratio.one_year:
            return 0
        if not classify:
            return long_term_debt_ratio.three_months - long_term_debt_ratio.one_year
        # 1 score if long term debt ratio is lower in the current year, else 0 score
        return int(long_term_debt_ratio.three_months < long_term_debt_ratio.one_year)

    def get_liquidity_score(self, fundamental, classify=True):
        '''Get the Liquidity score'''
        # if current or previous year's current ratio data does not exist, return 0 score
        current_ratio = fundamental.operation_ratios.current_ratio
        if not current_ratio.three_months or not current_ratio.one_year:
            return 0
        if not classify:
            return current_ratio.three_months - current_ratio.one_year
        # 1 score if current ratio is higher in the current year, else 0 score
        return int(current_ratio.three_months > current_ratio.one_year)

    def get_share_issued_score(self, fundamental, classify=True):
        '''Get the share issued score'''
        # if current or previous year's issued shares data does not exist, return 0 score
        shares_issued = fundamental.financial_statements.balance_sheet.share_issued
        if not shares_issued.three_months or not shares_issued.twelve_months:
            return 0
        if not classify:
            return shares_issued.three_months - shares_issued.twelve_months
        # 1 score if shares issued did not increase in the current year, else 0 score
        return int(shares_issued.three_months <= shares_issued.twelve_months)

    def get_gross_margin_score(self, fundamental, classify=True):
        '''Get the gross margin score'''
        # if current or previous year's gross margin data does not exist, return 0 score
        gross_margin = fundamental.operation_ratios.gross_margin
        if not gross_margin.three_months or not gross_margin.one_year:
            return 0
        if not classify:
            return gross_margin.three_months - gross_margin.one_year
        # 1 score if gross margin is higher in the current year, else 0 score
        return int(gross_margin.three_months > gross_margin.one_year)

    def get_asset_turnover_score(self, fundamental, classify=True):
        '''Get the asset turnover score'''
        # if current or previous year's asset turnover data does not exist, return 0 score
        asset_turnover = fundamental.operation_ratios.assets_turnover
        if not asset_turnover.three_months or not asset_turnover.one_year:
            return 0
        if not classify:
            return asset_turnover.three_months - asset_turnover.one_year
        # 1 score if asset turnover is higher in the current year, else 0 score
        return int(asset_turnover.three_months > asset_turnover.one_year)

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


class PiotroskiScoreUniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self, threshold, roa_multiple=1, universe_size=100):
        super().__init__(self._select_assets)
        self._piotroski_score = PiotroskiScore(roa_multiple)
        self._threshold = threshold
        self._universe_size = universe_size

    def _select_assets(self, fundamentals):
        # We use a dictionary to hold the F-Score of each stock
        f_scores = {
            f.symbol: self._piotroski_score.get_score(f) for f in fundamentals
            # We only want stocks with fundamental data and price > $1
            if f.has_fundamental_data and f.price > 1 and f.dollar_volume > 100_000
        }

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

        # Original Paper:
        # Select ALL stocks over the threshold.
        #return [symbol for symbol, fscore in f_scores.items() if fscore >= self._threshold]