#region imports
from AlgorithmImports import *

import statistics as stat
import pickle
from collections import deque
#endregion


class DynamicCalibratedGearbox(QCAlgorithm):

    def initialize(self):
        ### IMPORTANT: FOR USERS RUNNING THIS ALGORITHM IN LIVE TRADING,
        ### RUN THE BACKTEST ONCE
        
        self._tech_ROA_key = 'TECH_ROA'
        
        # we need 3 extra years to warmup our ROA values
        self.set_start_date(2016, 4, 1)  
        self.set_end_date(2020, 9, 1)
        self.set_warm_up(timedelta(3*365))
        
        self.set_cash(100000)  # Set Strategy Cash
        
        self.set_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(days=31)))
        self.set_execution(ImmediateExecutionModel())
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(lambda time:None))
        
        self.add_universe_selection(
            FineFundamentalUniverseSelectionModel(self._coarse_filter, self._fine_filter)
        )
        self.universe_settings.resolution = Resolution.DAILY
        
        self._curr_month = -1
        
        # store ROA of tech stocks
        self._tech_ROA = {}
        
        self._symbols = None
        
        if self.live_mode and not self.object_store.contains_key(self._tech_ROA_key):
            self.quit('QUITTING: USING LIVE MOVE WITHOUT TECH_ROA VALUES IN OBJECT STORE')
            
        self._quarters = 0
        
    def on_end_of_algorithm(self):
        self.log('Algorithm End')
        
        self._save_data()
        
    def _save_data(self):
        '''
        Saves the tech ROA data to ObjectStore
        '''
        
        # Symbol objects aren't picklable, hence why we use the ticker string
        tech_ROA = {symbol.value: roa for symbol, roa in self._tech_ROA.items()}
        self.object_store.save_bytes(self._tech_ROA_key, pickle.dumps(tech_ROA))
    
    def _coarse_filter(self, coarse):
        
        # load data from ObjectStore
        if len(self._tech_ROA) == 0 and self.object_store.contains_key(self._tech_ROA_key):
            tech_ROA = self.object_store.read_bytes(self._tech_ROA_key)
            tech_ROA = pickle.loads(bytearray(tech_ROA))
            self._tech_ROA = {Symbol.create(ticker, SecurityType.EQUITY, Market.USA): roa for ticker, roa in tech_ROA.items()}
            return list(self._tech_ROA.keys())

        if self._curr_month == self.time.month:
            return Universe.UNCHANGED
        
        self._curr_month = self.time.month
        
        # we only want to update our ROA values every three months
        if self.time.month % 3 != 1:
            return Universe.UNCHANGED
        
        self._quarters += 1
        
        return [c.symbol for c in coarse if c.has_fundamental_data]
      
    def _fine_filter(self, fine):
        # 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
        
        tech_securities = [f for f in fine if f.asset_classification.morningstar_sector_code == MorningstarSectorCode.TECHNOLOGY and
                                                not np.isnan(f.operation_ratios.roa.three_months)]
        
        for security in tech_securities:
            # we use deques instead of RWs since deques are picklable
            symbol = security.symbol
            if symbol not in self._tech_ROA:
                # 3 years * 4 quarters = 12 quarters of data
                self._tech_ROA[symbol] = deque(maxlen=12)
            self._tech_ROA[symbol].append(security.operation_ratios.roa.three_months)
        
        if self.live_mode:
            # this ensures we don't lose new data from an algorithm outage
            self.save_data()
        
        # 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 (self._quarters < 12 and not self.live_mode):
            return Universe.UNCHANGED
        
        # make sure our stocks has these fundamentals
        tech_securities = [
            x for x in tech_securities 
            if (
                not np.isnan(x.operation_ratios.roa.one_year) and
                x.operation_ratios.roa.one_year and
                not np.isnan(x.financial_statements.cash_flow_statement.operating_cash_flow.twelve_months) and
                x.financial_statements.cash_flow_statement.operating_cash_flow.twelve_months and
                not np.isnan(x.financial_statements.balance_sheet.total_assets.twelve_months) and
                x.financial_statements.balance_sheet.total_assets.twelve_months and
                not np.isnan(x.financial_statements.income_statement.research_and_development.twelve_months) and
                x.financial_statements.income_statement.research_and_development.twelve_months and
                not np.isnan(x.financial_statements.cash_flow_statement.cap_ex_reported.twelve_months) and
                x.financial_statements.cash_flow_statement.cap_ex_reported.twelve_months and
                not np.isnan(x.financial_statements.income_statement.selling_general_and_administration.twelve_months) and
                x.financial_statements.income_statement.selling_general_and_administration.twelve_months and
                not np.isnan(x.market_cap) and 
                x.market_cap
            )
        ]

        # compute the variance of the ROA for each tech stock
        tech_VARROA = {symbol:stat.variance(roa) for symbol, roa in self._tech_ROA.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 1-year
        tech_ROA1Y = {x.symbol:x.operation_ratios.roa.OneYear for x in tech_securities}
        tech_ROA1Y_median = stat.median(tech_ROA1Y.values())
        
        # Cash Flow ROA
        tech_CFROA = {x.symbol: (
            x.financial_statements.cash_flow_statement.operating_cash_flow.twelve_months 
            / x.financial_statements.balance_sheet.total_assets.twelve_months
            ) for x in tech_securities}
        tech_CFROA_median = stat.median(tech_CFROA.values())
        
        # R&D to MktCap
        tech_RD2MktCap = {x.symbol: (
            x.financial_statements.income_statement.research_and_development.twelve_months / x.market_cap
            ) for x in tech_securities}
        tech_RD2MktCap_median = stat.median(tech_RD2MktCap.values()) 
            
        # CapEx to MktCap
        tech_CaPex2MktCap = {x.symbol: (
            x.financial_statements.cash_flow_statement.cap_ex_reported.twelve_months / x.market_cap
            ) for x in tech_securities}
        tech_CaPex2MktCap_median = stat.median(tech_CaPex2MktCap.values())  
        
        # Advertising to MktCap
        tech_Ad2MktCap = {x.symbol: (
            x.financial_statements.income_statement.selling_general_and_administration.twelve_months / x.market_cap
            ) for x in tech_securities}
        tech_Ad2MktCap_median = stat.median(tech_Ad2MktCap.values())
        
        # sort fine by book-to-market ratio, get lower quintile
        has_book = [f for f in fine if f.financial_statements.balance_sheet.net_tangible_assets.twelve_months and f.market_cap]
        sorted_by_BM = sorted(has_book, key=lambda x: x.financial_statements.balance_sheet.net_tangible_assets.twelve_months / x.market_cap)[:len(has_book)//4]
        # choose tech stocks from lower quintile
        tech_symbols = [f.symbol for f in sorted_by_BM if f in tech_securities]
        
        ratioDicts_medians = [(tech_ROA1Y, tech_ROA1Y_median), 
                                (tech_CFROA, tech_CFROA_median), (tech_RD2MktCap, tech_RD2MktCap_median),
                                (tech_CaPex2MktCap, tech_CaPex2MktCap_median), (tech_Ad2MktCap, tech_Ad2MktCap_median)]
        
        def compute_g_score(symbol):
            g_score = 0
            if tech_CFROA[symbol] > tech_ROA1Y[symbol]:
                g_score += 1
            if symbol in tech_VARROA and tech_VARROA[symbol] < tech_VARROA_median:
                g_score += 1
            for ratio_dict, median in ratioDicts_medians:
                if symbol in ratio_dict and ratio_dict[symbol] > median:
                    g_score += 1
            return g_score
        
        # compute g-scores for each symbol    
        g_scores = {symbol:compute_g_score(symbol) for symbol in tech_symbols}

        return [symbol for symbol, g_score in g_scores.items() if g_score >= 5]