# region imports
from AlgorithmImports import *
# endregion
import statistics as stat
from collections import deque


class DynamicCalibratedGearbox(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.set_warm_up(timedelta(4*365))
        self.universe_settings.schedule.on(self.date_rules.month_start('SPY'))
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe_selection(FundamentalUniverseSelectionModel(self._select_assets))
        self.set_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(31)))
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(lambda time:None))
        self.set_execution(ImmediateExecutionModel())
        
        # store ROA of tech stocks
        self._roa_values_by_symbol = {}

    def _select_assets(self, fundamentals):
        # book value == FinancialStatements.BalanceSheet.NetTangibleAssets (book value and NTA are synonyms)
        # BM (Book-to-Market) == book value / MarketCap
        # ROA == OperationRatios.ROA
        # CFROA == FinancialStatements.CashFlowStatement.OperatingCashFlow / FinancialStatements.BalanceSheet.TotalAssets
        # R&D to MktCap == FinancialStatements.IncomeStatement.ResearchAndDevelopment / MarketCap
        # CapEx to MktCap == FinancialStatements.CashFlowStatement.CapExReported / MarketCap
        # Advertising to MktCap == FinancialStatements.IncomeStatement.SellingGeneralAndAdministration / MarketCap
        #   note: this parameter may be slightly higher than pure advertising costs

        # We only want to update our ROA values every three months.
        if self.time.month % 3 != 1:
            return Universe.UNCHANGED
        # Select tech stocks that have an ROA value.
        tech_securities = [
            f for f in fundamentals
            if (f.has_fundamental_data and
                f.asset_classification.morningstar_sector_code == MorningstarSectorCode.TECHNOLOGY and
                not np.isnan(f.operation_ratios.roa.three_months))
        ]
        # Record the ROA value of each stock.
        for f in tech_securities:
            symbol = f.symbol
            if symbol not in self._roa_values_by_symbol:
                # 3 years * 4 quarters = 12 quarters of data
                self._roa_values_by_symbol[symbol] = deque(maxlen=12)
            self._roa_values_by_symbol[symbol].append(f.operation_ratios.roa.three_months)
        # We want to rebalance in the fourth month after the (fiscal) year ends
        # so that we have the most recent quarter's data.
        if self.time.month != 4 or not any(len(ROA) == ROA.maxlen for ROA in self._roa_values_by_symbol.values()):
            return Universe.UNCHANGED
        # Make sure our stocks has these fundamentals.
        tech_securities = [
            f for f in tech_securities 
            if not any([np.isnan(factor) or factor == 0 for factor in [
                f.operation_ratios.roa.one_year,
                f.financial_statements.cash_flow_statement.operating_cash_flow.twelve_months,
                f.financial_statements.balance_sheet.total_assets.twelve_months,
                f.financial_statements.income_statement.research_and_development.twelve_months,
                f.financial_statements.cash_flow_statement.cap_ex_reported.twelve_months,
                f.financial_statements.income_statement.selling_general_and_administration.twelve_months,
                f.market_cap
            ]])
        ]                                               
        # Compute the variance of the ROA for each tech stock.
        tech_VARROA = {
            symbol : stat.variance(ROA) 
            for symbol, ROA in self._roa_values_by_symbol.items() if len(ROA) == ROA.maxlen
        }
        if len(tech_VARROA) < 2:
            return Universe.UNCHANGED
        tech_VARROA_median = stat.median(tech_VARROA.values())
        
        # We will now map tech Symbols to various fundamental ratios, 
        #   and compute the median for each ratio.
        roa_1y_by_symbol = {} # ROA 1-year
        cf_roa_by_symbol = {} # Cash Flow ROA
        rd_to_mkt_cap_by_symbol = {} # R&D to MktCap
        cap_ex_to_mkt_cap_by_symbol = {} # CapEx to MktCap
        ad_to_mkt_cap_by_symbol = {} # Advertising to MktCap
        for f in tech_securities:
            roa_1y_by_symbol[f.symbol] = f.operation_ratios.roa.one_year
            cf_roa_by_symbol[f.symbol] = (
                f.financial_statements.cash_flow_statement.operating_cash_flow.twelve_months
                / f.financial_statements.balance_sheet.total_assets.twelve_months
            ) 
            rd_to_mkt_cap_by_symbol[f.symbol] = (
                f.financial_statements.income_statement.research_and_development.twelve_months
                / f.market_cap
            ) 
            cap_ex_to_mkt_cap_by_symbol[f.symbol] = (
                f.financial_statements.cash_flow_statement.cap_ex_reported.twelve_months
                / f.market_cap
            ) 
            ad_to_mkt_cap_by_symbol[f.symbol] = (
                f.financial_statements.income_statement.selling_general_and_administration.twelve_months
                / f.market_cap
            )
        median_by_factors = [
            (factor_dict, stat.median(factor_dict.values()))
            for factor_dict in [
                roa_1y_by_symbol,
                cf_roa_by_symbol,
                rd_to_mkt_cap_by_symbol,
                cap_ex_to_mkt_cap_by_symbol,
                ad_to_mkt_cap_by_symbol,
            ]
        ]
        
        # Sort all stocks by book-to-market ratio and get the lower quintile.
        has_book = [
            f for f in fundamentals 
            if not any([np.isnan(factor) or factor == 0 for factor in [
                f.financial_statements.balance_sheet.net_tangible_assets.twelve_months,
                f.market_cap
            ]])
        ]
        sorted_by_BM = sorted(
            has_book, 
            key=lambda f: f.financial_statements.balance_sheet.net_tangible_assets.twelve_months / f.market_cap
        )[:len(has_book)//4]
        # Choose tech stocks from the lower quintile.
        tech_symbols = [f.symbol for f in sorted_by_BM if f in tech_securities]
        
        def compute_g_score(symbol):
            g_score = 0
            if cf_roa_by_symbol[symbol] > roa_1y_by_symbol[symbol]:
                g_score += 1
            if symbol in tech_VARROA and tech_VARROA[symbol] < tech_VARROA_median:
                g_score += 1
            for factor_by_symbol, median in median_by_factors:
                if symbol in factor_by_symbol and factor_by_symbol[symbol] > median:
                    g_score += 1
            return g_score
        
        # Compute g-scores for each asset.
        g_scores = {symbol : compute_g_score(symbol) for symbol in tech_symbols}
        # Apply G-score threshold.
        return [symbol for symbol, g_score in g_scores.items() if g_score >= 5]